diff --git a/.codespellrc b/.codespellrc index 306e5e4e..d8c810ec 100644 --- a/.codespellrc +++ b/.codespellrc @@ -1,3 +1,3 @@ [codespell] skip = *.lock,*.json,submodules/*,.venv/*,.git,docs/node_modules/* -ignore-words-list = dout,te,subtile,parm,mot,numer +ignore-words-list = dout,te,subtile,parm,mot,numer,notin diff --git a/.gitignore b/.gitignore index a6f1cb97..673ea19e 100644 --- a/.gitignore +++ b/.gitignore @@ -76,3 +76,13 @@ trace/ docs/node_modules/ docs/dist/ docs/.astro/ + +# Generated tokenized datasets and packing caches (per-benchmark) +experiments/local_benchmark/dataset_cache/ +experiments/local_benchmark/datasets/ + +# Kimi topology sweep scratch (tracked on branch `kimi-sweep-archive` instead) +experiments/local_benchmark/sweeps/ + +# Multi-adapter LoRA E2E artifacts can contain full per-rank logs and copied checkpoints. +experiments/multi_adapter_lora/results/ diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 38d2ae31..62af9fa4 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -4,7 +4,7 @@ 1. **Branch** off `main` with a descriptive name: `feature/my-feature`, `fix/bug-description` 2. **Commit** early and often on your branch — commit messages don't matter much here -3. **Open a PR** against `main` when ready for review +3. **Open a PR** against `main` when ready for review. The PR title must follow [Conventional Commits](https://www.conventionalcommits.org/) (see below), since it becomes the squash-merge commit message and drives automated release versioning. 4. **Squash merge** — all PRs are merged as a single squash commit; write a clean PR title and description since that becomes the commit message ## PR Guidelines @@ -12,7 +12,8 @@ - Keep PRs focused — one feature or fix per PR - Add tests for new behavior; existing tests must pass - Update relevant docs if behavior changes -- PR title should be imperative and descriptive: `Add chunked cross-entropy loss` not `chunked ce` +- PR title must follow Conventional Commits: `type: description` or `type(scope): description`. Allowed types: `feat`, `fix`, `perf`, `revert` (trigger releases); `chore`, `docs`, `test`, `refactor`, `ci`, `build` (no release). Append `!` for breaking changes (`feat!: ...` → major bump). An optional `[TICKET-123]` ticket prefix is allowed before the type. Examples: `feat: add chunked cross-entropy loss`, `fix(moe): correct expert routing`, `feat!: drop Python 3.9 support`. +- A CI check enforces this on every PR and comments the detected version bump. ## Commit Message (Squash Merge) diff --git a/README.md b/README.md index 6103e43b..63084028 100644 --- a/README.md +++ b/README.md @@ -22,7 +22,7 @@ XoRL is a distributed training framework designed for large language models with | Repo                        | Description | |---|---| -| **[xorl](https://github.com/togethercomputer/xorl-internal)** | Distributed training framework — local SFT/pretraining and server-mode RL training | +| **[xorl](https://github.com/togethercomputer/xorl)** | 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 | @@ -48,8 +48,8 @@ XoRL is a distributed training framework designed for large language models with ## 🚀 Installation ```bash -git clone --recurse-submodules git@github.com:togethercomputer/xorl-internal.git -cd xorl-internal +git clone --recurse-submodules git@github.com:togethercomputer/xorl.git +cd xorl ``` > Already cloned without `--recurse-submodules`? Run `git submodule update --init --recursive` @@ -102,7 +102,7 @@ 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). +See the [installation guide](https://togethercomputer.github.io/xorl/getting-started/installation/) for full setup including optional dependencies (DeepEP, Flash Attention). ## ⚡ Quick Start diff --git a/docs/src/content/docs/config-reference/local.md b/docs/src/content/docs/config-reference/local.md index 5bead0b1..aa9643d0 100644 --- a/docs/src/content/docs/config-reference/local.md +++ b/docs/src/content/docs/config-reference/local.md @@ -35,7 +35,7 @@ torchrun --nproc_per_node=8 -m xorl.cli.train config.yaml \ | `ep_dispatch` | `alltoall` | Expert-parallel dispatch: `alltoall` or `deepep` (NVLink-optimized). | | `deepep_buffer_size_gb` | `2.0` | DeepEP NVLink buffer size per GPU in GB. Only active when `ep_dispatch: deepep`. | | `deepep_num_sms` | `20` | SMs assigned to DeepEP communication kernels. Must be even. Lower values leave more SMs for overlapped compute. | -| `deepep_async_combine` | `false` | Overlap DeepEP combine with the next layer's compute (experimental). | +| `deepep_async_combine` | `false` | Overlap DeepEP combine with the next layer's compute (experimental, unsafe). Forced to `false` in code unless `XORL_DEEPEP_UNSAFE_ASYNC_COMBINE=1` is exported; without that env var, deferring the comm-stream sync races the transformer block's read of the combined tensor on the default stream. | | `merge_qkv` | `true` | Keep Q/K/V projections fused as `qkv_proj`. Set `false` for tensor parallelism or per-projection LoRA. | | `basic_modules` | `[]` | Additional module names (beyond `_no_split_modules`) to shard as separate FSDP units. | | `foundation` | `{}` | Extra foundation model config (dict). | @@ -167,7 +167,7 @@ Each entry in `datasets` (or `test_datasets`) is a dict: | `activation_gpu_limit` | `0.0` | GB of activations to keep on GPU when offloading. `0.0` = offload all. | | `enable_compile` | `false` | `torch.compile` for model forward pass. | | `init_device` | `cuda` | Device for weight initialization: `cpu` (rank 0 only), `cuda`, `meta` (required for FSDP2), `npu`. | -| `load_weights_mode` | `broadcast` | `broadcast`: rank 0 reads weights, broadcasts to other ranks (reduces disk I/O). `all_ranks`: every rank reads from disk. | +| `load_weights_mode` | `grouped` | `grouped`: one reader per node for dense/shared weights plus one reader per EP-FSDP group for expert weights, with rank-0 fallback when grouped fanout groups are unavailable. `all_ranks`: every rank reads from disk. `skip`: skip HuggingFace weight loading and materialize model weights from `load_checkpoint_path` (DCP). | | `enable_full_determinism` | `false` | Full determinism mode. Requires `allow_cuda_launch_blocking: true`. Degrades performance. | | `allow_cuda_launch_blocking` | `false` | Allow `CUDA_LAUNCH_BLOCKING=1`. Off by default to prevent accidental performance degradation. | | `empty_cache_steps` | `500` | Call `torch.cuda.empty_cache()` every N steps. | diff --git a/docs/src/content/docs/config-reference/server.md b/docs/src/content/docs/config-reference/server.md index 21eccae1..876273c9 100644 --- a/docs/src/content/docs/config-reference/server.md +++ b/docs/src/content/docs/config-reference/server.md @@ -33,7 +33,7 @@ python -m xorl.server.launcher --mode auto --config config.yaml \ | `ep_dispatch` | `alltoall` | Expert-parallel dispatch: `alltoall` or `deepep` (NVLink-optimized). | | `deepep_buffer_size_gb` | `2.0` | DeepEP NVLink buffer size per GPU in GB. Only active when `ep_dispatch: deepep`. | | `deepep_num_sms` | `20` | SMs assigned to DeepEP communication kernels. Must be even. | -| `deepep_async_combine` | `false` | Overlap DeepEP combine with the next layer's compute (experimental). | +| `deepep_async_combine` | `false` | Overlap DeepEP combine with the next layer's compute (experimental, unsafe). Forced to `false` in code unless `XORL_DEEPEP_UNSAFE_ASYNC_COMBINE=1` is exported; without that env var, deferring the comm-stream sync races the transformer block's read of the combined tensor on the default stream. | | `merge_qkv` | `true` | Keep Q/K/V projections fused. Set `false` for tensor parallelism. | | `basic_modules` | `[]` | Additional module names to shard as separate FSDP units. | | `foundation` | `{}` | Foundation model extra config (dict). | @@ -86,7 +86,7 @@ These flags align the training model's numerics with the inference engine (SGLan | `enable_reentrant` | `false` | Use reentrant gradient checkpointing. | | `enable_forward_prefetch` | `false` | FSDP forward prefetch. | | `init_device` | `meta` | Model initialization device: `cpu`, `meta`, `cuda`. | -| `load_weights_mode` | `auto` | Weight loading: `auto`, `safetensors`, `dcp`. | +| `load_weights_mode` | `grouped` | Weight loading mode: `grouped` (default, with rank-0 fallback), `all_ranks`, or `skip`. | | `ce_mode` | `compiled` | Cross-entropy implementation: `compiled` (recommended, `torch.compile`) or `eager` (may OOM at 32K+ seq len). | --- @@ -170,6 +170,7 @@ ZMQ communication between the launcher, workers, and API server. | `qlora_exclude_modules` | `null` | Modules to exclude from quantization (e.g., `[lm_head]`). | | `merge_lora_interval` | `0` | Merge LoRA into base weights every N steps. `0` = never. | | `reset_optimizer_on_merge` | `false` | ReLoRA optimizer reset after merge. | +| `adapter_state_load_mode` | `all_ranks` | How to restore multi-adapter checkpoints: `all_ranks` loads on every rank; `rank0_broadcast` loads on rank 0 and broadcasts weights, metadata, and optimizer state. | --- diff --git a/docs/src/content/docs/getting-started/installation.md b/docs/src/content/docs/getting-started/installation.md index c9e37a6a..02c80ab9 100644 --- a/docs/src/content/docs/getting-started/installation.md +++ b/docs/src/content/docs/getting-started/installation.md @@ -13,8 +13,8 @@ title: "Installation" ## Clone the repo ```bash -git clone --recurse-submodules https://github.com/togethercomputer/xorl-internal -cd xorl-internal +git clone --recurse-submodules https://github.com/togethercomputer/xorl +cd xorl ``` > Already cloned without `--recurse-submodules`? Run `git submodule update --init --recursive` diff --git a/docs/src/content/docs/getting-started/quickstart.md b/docs/src/content/docs/getting-started/quickstart.md index f97288a9..e0abed7b 100644 --- a/docs/src/content/docs/getting-started/quickstart.md +++ b/docs/src/content/docs/getting-started/quickstart.md @@ -109,7 +109,7 @@ future = requests.post(f"{base_url}/api/v1/optim_step", 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`. +[`examples/server/no_robot_sft/`](https://github.com/togethercomputer/xorl/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 @@ -131,7 +131,7 @@ The script uses `xorl_client.TrainingClient` to drive a LoRA SFT loop with onlin ### 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. +[`examples/server/password_memorization/`](https://github.com/togethercomputer/xorl/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 @@ -149,7 +149,7 @@ 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. +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/tree/main/examples/server/password_memorization/README.md) for the full test matrix across Qwen3-8B, Qwen3-30B, and Qwen3-235B. ## LoRA Fine-tuning diff --git a/docs/src/content/docs/moe/deepep.mdx b/docs/src/content/docs/moe/deepep.mdx index 51b2e196..d2d5e714 100644 --- a/docs/src/content/docs/moe/deepep.mdx +++ b/docs/src/content/docs/moe/deepep.mdx @@ -88,7 +88,7 @@ model: | `ep_dispatch` | `alltoall` | Set to `deepep` to enable | | `deepep_buffer_size_gb` | `2.0` | Per-GPU NVLink buffer pool in GB. Larger = fewer chunked transfers. Rule of thumb: `2 × token_budget × hidden_dim × sizeof(bf16)` | | `deepep_num_sms` | `20` | SMs dedicated to communication kernels. Must be even. | -| `deepep_async_combine` | `false` | Overlap combine with next layer's compute (experimental) | +| `deepep_async_combine` | `false` | Overlap combine with next layer's compute (experimental, **currently disabled** — see [Async Combine](#async-combine-experimental)) | --- @@ -223,17 +223,29 @@ If you see OOM during initialization, reduce `deepep_buffer_size_gb`. If profili ## Async Combine (Experimental) -`deepep_async_combine: true` overlaps the combine communication (outputs flowing back from expert ranks) with the next layer's compute: +:::caution[Currently disabled] +The async combine path is **forced off in code**. The combined +tensor is consumed by the rest of the transformer block (residual add, norm, +attention) on the default stream **before** the next `token_pre_dispatch()` +runs, so deferring `event.current_stream_wait()` races those consumers. +Setting `deepep_async_combine: true` alone is a no-op (logs a one-time warning). +To actually run async, also export `XORL_DEEPEP_UNSAFE_ASYNC_COMBINE=1` — only +do this if you have separately ensured the downstream consumer waits on the +combine event. +::: + +`deepep_async_combine: true` requests overlap of the combine communication +(outputs flowing back from expert ranks) with the next layer's compute: ``` Step N: dispatch → compute → [combine starts] Step N+1: [combine finishes] + next layer compute (overlapped) ``` -**Benefit:** Hides combine latency behind useful compute, especially when combine > dispatch (typical for large output projections). +**Benefit (when working):** Hides combine latency behind useful compute, especially when combine > dispatch (typical for large output projections). **Limitations:** -- Experimental — correctness verified on Qwen3 but not all architectures +- Gated off behind `XORL_DEEPEP_UNSAFE_ASYNC_COMBINE=1` - Requires careful ordering of CUDA streams - Not compatible with pipeline parallelism (PP > 1) @@ -247,7 +259,7 @@ On Qwen3-235B-A22B with EP=64 (8 nodes, 64 H100 NVLink GPUs): |---|---|---| | AllToAll (NCCL) | ~5s per forward_backward | Baseline | | DeepEP (num_sms=20) | ~1s per forward_backward | ~5× faster | -| DeepEP + async_combine | ~0.6s per forward_backward | ~8× faster | +| DeepEP + async_combine | ~0.6s per forward_backward | ~8× faster (currently gated off, see [Async Combine](#async-combine-experimental)) | Gains are most pronounced at high EP sizes (≥ 32) where AllToAll's O(EP) staging bottleneck dominates. @@ -259,7 +271,7 @@ Gains are most pronounced at high EP sizes (≥ 32) where AllToAll's O(EP) stagi | EP ≤ 8, multi-node InfiniBand | AllToAll — IB already uses RDMA | | EP ≥ 16, NVLink cluster | DeepEP | | EP ≥ 32, NVLink cluster | DeepEP strongly recommended | -| EP ≥ 64, NVLink cluster | DeepEP + async_combine | +| EP ≥ 64, NVLink cluster | DeepEP (async_combine currently gated off) | --- diff --git a/docs/src/content/docs/parallelism/data_parallelism.mdx b/docs/src/content/docs/parallelism/data_parallelism.mdx index 02d105fe..630caa6a 100644 --- a/docs/src/content/docs/parallelism/data_parallelism.mdx +++ b/docs/src/content/docs/parallelism/data_parallelism.mdx @@ -205,7 +205,7 @@ train: enable_full_shard: true enable_mixed_precision: true init_device: meta # required for FSDP2 - load_weights_mode: broadcast # rank0 loads, broadcasts to all ranks + load_weights_mode: grouped # grouped fanout, with rank-0 fallback ``` `init_device: meta` is **required** for FSDP2. Parameters are initially created on the meta device (zero-cost), then materialized by FSDP2 after `fully_shard` is applied. @@ -267,7 +267,7 @@ train: enable_full_shard: true enable_mixed_precision: true init_device: meta - load_weights_mode: broadcast + load_weights_mode: grouped ``` The constraint that must hold: `data_parallel_shard_size × data_parallel_replicate_size == data_parallel_size` (where `data_parallel_size = world_size / (TP × PP × CP)`). xorl validates this in `TrainingArguments.__post_init__`. @@ -478,7 +478,7 @@ train: enable_gradient_checkpointing: true enable_full_shard: true init_device: meta - load_weights_mode: broadcast + load_weights_mode: grouped ``` **Memory profile (Qwen3-8B, ~8B params):** @@ -505,7 +505,7 @@ train: enable_gradient_checkpointing: true enable_full_shard: true init_device: meta - load_weights_mode: broadcast + load_weights_mode: grouped ``` Here `world_size = dp_shard × ulysses = 2 × 4 = 8`. Each 2-GPU FSDP shard group uses 4-way Ulysses to handle long sequences that would not fit on 2 GPUs individually. @@ -522,7 +522,7 @@ train: enable_gradient_checkpointing: true enable_full_shard: true init_device: meta - load_weights_mode: broadcast + load_weights_mode: grouped ``` Cross-node IB traffic is limited to the all-reduce of *averaged gradients* (not full shard all-gathers), which is typically 4–8x less than pure FSDP2 would require across nodes. @@ -578,7 +578,7 @@ train: enable_full_shard: true reshard_after_forward: false # PP: keep params gathered between fwd micro-batches init_device: meta - load_weights_mode: broadcast + load_weights_mode: grouped ``` `world_size = PP × dp_shard = 2 × 4 = 8`. Each PP stage has 4 GPUs running FSDP2 over non-expert params and EP=4 over expert params. @@ -643,15 +643,16 @@ When EP is enabled, xorl replaces FSDP2's automatic prefetching with explicit pe |-------|-------------|-----------------| | `meta` | Parameters on meta device; materialized lazily by FSDP2 | FSDP2 only (required) | | `cuda` | Parameters initialized directly on GPU | DDP, or FSDP2 debugging | -| `cpu` | Parameters on CPU (rank 0 only for broadcast) | DDP only; not supported with EP | +| `cpu` | Parameters on CPU (rank 0 only for rank-0 fallback) | DDP only; not supported with EP | | `npu` | Ascend NPU device | DDP, FSDP2 | ### `load_weights_mode` | Value | Description | When to use | |-------|-------------|-------------| -| `broadcast` (default) | Rank 0 reads checkpoint from disk, broadcasts shards to all ranks | Default; avoids N-way disk I/O bottleneck | +| `grouped` (default) | Dense/shared tensors are read once per node and expert tensors once per EP-FSDP group, then replicated/scattered within those groups | Default; best for large MoE/EP loads and falls back to rank-0 loading when grouped fanout groups are unavailable | | `all_ranks` | Every rank reads the checkpoint independently | Fast parallel storage (e.g., Lustre, object storage), or when EP requires each rank to load its own expert shard | +| `skip` | Skip HuggingFace checkpoint loading and load model weights later from `load_checkpoint_path` | Model-only DCP checkpoints created by `scripts/convert_checkpoint.py --save-format dcp` | --- @@ -665,9 +666,9 @@ Meta-device initialization creates parameter tensors with zero CPU/GPU memory co init_device: meta ``` -### Use `load_weights_mode: broadcast` by default +### Use `load_weights_mode: grouped` by default -`broadcast` mode has rank 0 load the checkpoint from disk and distribute shards. This is the safest option when the storage system cannot handle parallel reads from all ranks simultaneously. Switch to `all_ranks` only on parallel filesystems with guaranteed per-rank I/O performance, or when loading EP shards that differ per rank. +`grouped` mode uses grouped fanout for large MoE checkpoints so dense/shared weights are read once per node and expert weights are read once per EP-FSDP group. When grouped fanout groups are unavailable, it falls back to rank-0 loading. Switch to `all_ranks` only on parallel filesystems with guaranteed per-rank I/O performance. ### Set `reshard_after_forward: false` for pipeline parallelism diff --git a/docs/src/content/docs/parallelism/expert_parallelism.mdx b/docs/src/content/docs/parallelism/expert_parallelism.mdx index 92deafed..eabc55af 100644 --- a/docs/src/content/docs/parallelism/expert_parallelism.mdx +++ b/docs/src/content/docs/parallelism/expert_parallelism.mdx @@ -176,7 +176,7 @@ parameters, matches each expert parameter against the EP plan, then either: - **Redistributes** the full tensor into a local shard using `DTensor.redistribute()` with a `Shard(0)` placement on the EP mesh (normal path when loading from a checkpoint). - **Annotates** already-local tensors with a `spec_info` attribute (fast path when weights - were loaded EP-aware, e.g. with `load_weights_mode: all_ranks`). + were loaded EP-aware, e.g. with `load_weights_mode: all_ranks` or `grouped`). After sharding, each rank's expert parameter has shape `[num_local_experts, ...]`. The `spec_info` attribute records the `ep_fsdp_mesh`, the shard placement, and the FQN for @@ -357,9 +357,20 @@ The forward dispatch is a single `torch.autograd.Function` boundary: ### Async combine overlap -Setting `deepep_async_combine: true` enables the async combine path. The combine -communication for layer `L` overlaps with the attention/dense-FFN computation of layer -`L+1`. The pending event is synchronized at the start of the next layer's dispatch. +Setting `deepep_async_combine: true` requests the async combine path, where the +combine communication for layer `L` would overlap with the attention/dense-FFN +computation of layer `L+1` and the pending event would be synchronized at the +start of the next layer's dispatch. + +:::caution +This path is currently **disabled at the kernel boundary** because the combined +tensor is read by the rest of the transformer block (residual add, norm, attention) +on the default stream **before** the next `token_pre_dispatch()` runs, so deferring +the `current_stream_wait` races those consumers. `tokens_post_combine` forces +`async_combine=False` unless `XORL_DEEPEP_UNSAFE_ASYNC_COMBINE=1` is exported. +Setting the config flag alone is a no-op (with a one-time warning); set the env +var only if you have separately guaranteed the consumer waits on the combine event. +::: ### Supported EP sizes @@ -826,7 +837,7 @@ ranks. | `ep_dispatch` | `str` | `"alltoall"` | EP token dispatch strategy: `"alltoall"` or `"deepep"`. | | `deepep_buffer_size_gb` | `float` | `2.0` | NVLink staging buffer size in GB for DeepEP. | | `deepep_num_sms` | `int` | `20` | SMs allocated to DeepEP communication kernels (must be even). Lower values leave more for expert compute. | -| `deepep_async_combine` | `bool` | `False` | Overlap DeepEP combine with next layer's compute. | +| `deepep_async_combine` | `bool` | `False` | Overlap DeepEP combine with next layer's compute. Currently gated off in code; requires `XORL_DEEPEP_UNSAFE_ASYNC_COMBINE=1` (see [Async combine overlap](#async-combine-overlap)). | ### Train config block @@ -837,7 +848,7 @@ ranks. | `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`. | -| `load_weights_mode` | `str` | `"broadcast"` | `"all_ranks"` lets every rank read its local expert shard directly without a broadcast. Preferred for EP. | +| `load_weights_mode` | `str` | `"grouped"` | Default grouped fanout: one dense/shared reader per node and one expert reader per EP-FSDP group, with rank-0 fallback when grouped fanout groups are unavailable. Use `"all_ranks"` on storage that can handle one checkpoint read per rank, or `"skip"` with `load_checkpoint_path` to load model weights from a converted DCP checkpoint. | ### Environment variables diff --git a/docs/src/content/docs/parallelism/pipeline_parallelism.mdx b/docs/src/content/docs/parallelism/pipeline_parallelism.mdx index f8ddead5..927f40ea 100644 --- a/docs/src/content/docs/parallelism/pipeline_parallelism.mdx +++ b/docs/src/content/docs/parallelism/pipeline_parallelism.mdx @@ -696,7 +696,7 @@ train: enable_gradient_checkpointing: true enable_full_shard: true init_device: meta - load_weights_mode: broadcast + load_weights_mode: grouped ``` **Process group layout** (8 GPUs, ranks 0–7): diff --git a/docs/src/content/docs/parallelism/tensor_parallelism.mdx b/docs/src/content/docs/parallelism/tensor_parallelism.mdx index b5c73869..1397f31f 100644 --- a/docs/src/content/docs/parallelism/tensor_parallelism.mdx +++ b/docs/src/content/docs/parallelism/tensor_parallelism.mdx @@ -411,7 +411,7 @@ When `init_device: meta` is used, weights must be loaded **after** TP is applied if parallel_state.tp_enabled: model = parallelize_module(model, ...) # apply TP first if kwargs.get("init_device") == "meta" and weights_path is not None: - rank0_load_and_broadcast_weights(...) # load weights into TP-sharded params + grouped_load_weights(...) # load weights into TP-sharded params kwargs["skip_weight_loading"] = True # skip again in FSDP path # then FSDP2 wraps the already-materialized TP DTensors @@ -462,7 +462,7 @@ train: enable_gradient_checkpointing: true enable_compile: true # recommended with TP for fused kernels init_device: meta - load_weights_mode: broadcast + load_weights_mode: grouped ``` GPU layout (TP groups across columns, FSDP shard groups across rows): @@ -504,7 +504,7 @@ train: enable_mixed_precision: true enable_gradient_checkpointing: true init_device: meta - load_weights_mode: broadcast + load_weights_mode: grouped ``` ### TP=4 + compile (from `qwen3_8b_tp4_compile.yaml`) @@ -529,7 +529,7 @@ train: enable_gradient_checkpointing: true enable_compile: true init_device: meta - load_weights_mode: broadcast + load_weights_mode: grouped ``` `torch.compile` is applied to each decoder block before FSDP wrapping (`build_parallelize_model` handles the ordering). Compiled kernels can fuse the TP linear operations with activation functions for better throughput. @@ -615,7 +615,7 @@ For the same memory budget, FSDP2 with larger shard groups is generally more eff | `pipeline_parallel_size` | `int` | `1` | PP stage count. Compose with TP by assigning `tensor_parallel_size × pipeline_parallel_size` GPUs to model parallelism. | | `data_parallel_mode` | `str` | `"fsdp2"` | Must be `"fsdp2"` when using TP with meta-device initialization. | | `init_device` | `str` | `"meta"` | Use `"meta"` with TP to avoid materializing full weights before TP sharding. | -| `load_weights_mode` | `str` | `"broadcast"` | With TP, rank 0 loads from disk and broadcasts to TP peers. Use `"all_ranks"` if filesystem supports parallel reads. | +| `load_weights_mode` | `str` | `"grouped"` | Default grouped fanout, with rank-0 fallback when grouped fanout groups are unavailable. Use `"all_ranks"` if filesystem supports parallel reads, or `"skip"` with `load_checkpoint_path` for converted DCP checkpoints. | | `enable_compile` | `bool` | `false` | `torch.compile` per decoder block. Recommended with TP for fused kernels; must be applied before FSDP wrapping (handled automatically). | ### Model arguments (`src/xorl/arguments.py`, `ModelArguments` dataclass) @@ -663,7 +663,7 @@ def _resolve_tp_style(style_str): 1. Build model on meta device (no weights allocated) 2. Call unfuse_for_tp() — splits fused projections (still meta) 3. Call parallelize_module() — annotates params as DTensors with TP placement (still meta) -4. Load weights (rank0_load_and_broadcast_weights) — materializes meta DTensors into TP-sharded real tensors +4. Load weights (`grouped`) — materializes meta DTensors into TP-sharded real tensors 5. Call parallelize_model_fsdp2() — wraps TP DTensors with FSDP2 ``` 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 index 0e0d40d3..6526dea5 100644 --- a/docs/src/content/docs/server-training/client-sdk/loss-functions.md +++ b/docs/src/content/docs/server-training/client-sdk/loss-functions.md @@ -27,7 +27,7 @@ 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)): +Full PPO-style clipped policy gradient loss ([source](https://github.com/togethercomputer/xorl/blob/main/src/xorl/ops/loss/policy_loss.py)): ``` ratio = exp(new_logprobs - old_logprobs) @@ -65,7 +65,7 @@ fwd = client.forward_backward(data, loss_fn="policy_loss", loss_fn_params={ ## 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)): +Simpler importance-sampling loss ([source](https://github.com/togethercomputer/xorl/blob/main/src/xorl/ops/loss/importance_sampling_loss.py)): ``` ratio = exp(new_logprobs - old_logprobs) diff --git a/docs/src/content/docs/server-training/examples/password-memorization.md b/docs/src/content/docs/server-training/examples/password-memorization.md index e569aa47..6c0e72ea 100644 --- a/docs/src/content/docs/server-training/examples/password-memorization.md +++ b/docs/src/content/docs/server-training/examples/password-memorization.md @@ -2,7 +2,7 @@ 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. +[`examples/server/password_memorization/`](https://github.com/togethercomputer/xorl/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:** @@ -70,4 +70,4 @@ python examples/server/password_memorization/run_password_test.py \ | 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. +See the [example README](https://github.com/togethercomputer/xorl/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 index 9f61aef4..3777f914 100644 --- a/docs/src/content/docs/server-training/examples/sft-no-robots.md +++ b/docs/src/content/docs/server-training/examples/sft-no-robots.md @@ -2,7 +2,7 @@ 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. +[`examples/server/no_robot_sft/`](https://github.com/togethercomputer/xorl/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` diff --git a/docs/src/content/docs/server-training/overview.md b/docs/src/content/docs/server-training/overview.md index e932590f..60ca7caf 100644 --- a/docs/src/content/docs/server-training/overview.md +++ b/docs/src/content/docs/server-training/overview.md @@ -108,6 +108,8 @@ For multi-node setup, server configuration, and launcher CLI options, see [Launc The [xorl-client](/xorl/server-training/client-sdk/overview/) Python SDK drives the training server — see the [Client SDK page](/xorl/server-training/client-sdk/overview/) for installation, client classes, loss functions, and training loop examples. +Multi-adapter LoRA server training is multi-tenant: each `model_id` owns a session-specific LoRA and optimizer runtime spec, and non-default sessions can be removed with `kill_session`. Full-weight server training remains single-tenant and only supports the reserved `default` session. + --- ## Tinker API Compatibility diff --git a/docs/src/content/docs/server-training/training-server/api-reference.md b/docs/src/content/docs/server-training/training-server/api-reference.md index fb299587..b566f9ac 100644 --- a/docs/src/content/docs/server-training/training-server/api-reference.md +++ b/docs/src/content/docs/server-training/training-server/api-reference.md @@ -24,9 +24,9 @@ All endpoints are served at `http://:/`. Training operations use a t | Method | Path | Description | |---|---|---| -| `POST` | `/api/v1/create_model` | Create and register a new model session (LoRA or full-weight). | +| `POST` | `/api/v1/create_model` | Create and register a new training session. LoRA mode supports multi-tenant sessions; full-weight mode only supports the reserved `model_id="default"` session. | | `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. | +| `POST` | `/api/v1/kill_session` | Kill an active session. In LoRA mode, non-default tenant sessions are removed; in full-weight mode, the single active session is reset. | | `GET` | `/api/v1/session_info` | List active sessions and their state. | | `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. | diff --git a/docs/src/content/docs/training/local_training.mdx b/docs/src/content/docs/training/local_training.mdx index ba602512..b9eb7406 100644 --- a/docs/src/content/docs/training/local_training.mdx +++ b/docs/src/content/docs/training/local_training.mdx @@ -132,7 +132,7 @@ train: enable_gradient_checkpointing: true enable_full_shard: true init_device: meta - load_weights_mode: broadcast + load_weights_mode: grouped save_steps: 500 ckpt_manager: dcp ``` @@ -153,10 +153,10 @@ xorl downloads from HF Hub on first use (cached in `~/.cache/huggingface`). Pass ```yaml train: init_device: meta - load_weights_mode: broadcast # rank 0 reads, all ranks receive via NCCL + load_weights_mode: grouped # grouped fanout, with rank-0 fallback ``` -Meta device creates model parameters as zero-cost placeholders. FSDP2 then loads and shards weights directly into each rank's shard — no rank ever holds the full model in memory. Without meta device, every rank would briefly allocate the entire model before sharding, causing OOM for models larger than ~30B parameters on 80 GB GPUs. Use `load_weights_mode: all_ranks` for NVMe-fast local storage (each rank reads independently). +Meta device creates model parameters as zero-cost placeholders. FSDP2 then loads and shards weights directly into each rank's shard — no rank ever holds the full model in memory. Without meta device, every rank would briefly allocate the entire model before sharding, causing OOM for models larger than ~30B parameters on 80 GB GPUs. `grouped` is the default; it uses grouped fanout for large MoE/EP loads and falls back to rank-0 loading when grouped fanout groups are unavailable. Use `load_weights_mode: all_ranks` for NVMe-fast local storage (each rank reads independently). ## Data Parallel Modes diff --git a/docs/src/content/docs/training/system.mdx b/docs/src/content/docs/training/system.mdx index dec69ef8..3124472b 100644 --- a/docs/src/content/docs/training/system.mdx +++ b/docs/src/content/docs/training/system.mdx @@ -70,7 +70,7 @@ This page covers how `Trainer` is structured for local training — the training |---|---|---| | 1 | `_bootstrap()` | Init distributed (torchrun env), set device, seed, build `ParallelState` | | 2 | `_build_model()` | Load foundation model config; inject LoRA/QLoRA if enabled | -| 3 | `_parallelize()` | Apply TP plan, FSDP2/EP mesh, PP split; broadcast weights from rank 0 | +| 3 | `_parallelize()` | Apply TP plan, FSDP2/EP mesh, PP split; load weights with grouped fanout | | 4 | `_build_data()` | Load tokenizer, prepare dataset, build `DataLoaderBuilder` with collator chain | | 5 | `_build_optimizer()` | Create AdamW/Muon optimizer; build LR scheduler | | 6 | `_setup_observability()` | Configure structured logging, W&B | @@ -275,7 +275,7 @@ Packed bins are cached to disk (keyed by a hash of the packing config) so subseq 4. Non-expert FSDP — fully_shard(layer, mesh=fsdp_mesh) per decoder block 5. Root FSDP — fully_shard(model, mesh=fsdp_mesh) embeddings + lm_head 6. PP split — pipeline_module_split() → build_pp_stage() per rank -7. Weight loading — broadcast from rank 0 (or all_ranks read independently) +7. Weight loading — grouped fanout (or all_ranks read independently) ``` For TP, the plan is read from the model config's `base_model_tp_plan` dict (maps FQN patterns → `colwise`/`rowwise`/`embedding` style strings). xorl resolves these to PyTorch `ParallelStyle` objects and calls `parallelize_module()`. diff --git a/examples/local/coderforge/configs/qwen3_5_35b_a3b.yaml b/examples/local/coderforge/configs/qwen3_5_35b_a3b.yaml index 2e781c14..8b6c31f5 100644 --- a/examples/local/coderforge/configs/qwen3_5_35b_a3b.yaml +++ b/examples/local/coderforge/configs/qwen3_5_35b_a3b.yaml @@ -46,7 +46,7 @@ train: enable_full_shard: true enable_activation_offload: false init_device: meta - load_weights_mode: broadcast + load_weights_mode: grouped enable_full_determinism: false empty_cache_steps: 500 ckpt_manager: dcp diff --git a/examples/local/coderforge/configs/qwen3_5_35b_a3b_muon.yaml b/examples/local/coderforge/configs/qwen3_5_35b_a3b_muon.yaml index 4dcc0cd1..733428c1 100644 --- a/examples/local/coderforge/configs/qwen3_5_35b_a3b_muon.yaml +++ b/examples/local/coderforge/configs/qwen3_5_35b_a3b_muon.yaml @@ -52,7 +52,7 @@ train: enable_full_shard: true enable_activation_offload: false init_device: meta - load_weights_mode: broadcast + load_weights_mode: grouped enable_full_determinism: false empty_cache_steps: 500 ckpt_manager: dcp diff --git a/examples/local/dummy/configs/full/glm4_moe_cp1_sp8.yaml b/examples/local/dummy/configs/full/glm4_moe_cp1_sp8.yaml new file mode 100644 index 00000000..cbda872a --- /dev/null +++ b/examples/local/dummy/configs/full/glm4_moe_cp1_sp8.yaml @@ -0,0 +1,56 @@ +## GLM-4 MoE + Pure Ulysses: ep=8, cp=1, sp=8, dp=1 +model: + model_path: zai-org/GLM-4.5-Air + attn_implementation: flash_attention_3 + moe_implementation: triton + +data: + datasets: + - path: dummy + type: tokenized + max_seq_len: 16001 + select_columns: [input_ids, labels] + dataset_prepared_path: last_prepared_dataset + sample_packing_method: sequential + sample_packing_sequence_len: 128000 + dataloader_num_workers: 4 + dataloader_prefetch_factor: 2 + dataloader_pin_memory: true + +train: + output_dir: outputs/bench-glm4-moe-cp1-sp8 + data_parallel_mode: fsdp2 + ulysses_parallel_size: 2 + ringattn_parallel_size: 1 + expert_parallel_size: 8 + data_parallel_replicate_size: 1 + data_parallel_shard_size: 4 + + num_train_epochs: 5 + max_steps: 3 + + micro_batch_size: 1 + gradient_accumulation_steps: 1 + empty_cache_steps: 100 + + 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: all_ranks + enable_full_determinism: false + + ckpt_manager: dcp + save_steps: 0 + save_hf_weights: false + + use_wandb: false diff --git a/examples/local/dummy/configs/full/glm4_moe_ep8.yaml b/examples/local/dummy/configs/full/glm4_moe_ep8.yaml new file mode 100644 index 00000000..122d8018 --- /dev/null +++ b/examples/local/dummy/configs/full/glm4_moe_ep8.yaml @@ -0,0 +1,54 @@ +model: + model_path: zai-org/GLM-4.5-Air + attn_implementation: flash_attention_3 + moe_implementation: triton + +data: + datasets: + - path: dummy + type: tokenized + max_seq_len: 2048 + select_columns: [input_ids, labels] + dataset_prepared_path: last_prepared_dataset + sample_packing_method: sequential + sample_packing_sequence_len: 8192 + dataloader_num_workers: 4 + dataloader_prefetch_factor: 2 + dataloader_pin_memory: true + +train: + output_dir: outputs/bench-glm4-moe-ep8 + data_parallel_mode: fsdp2 + tensor_parallel_size: 1 + ulysses_parallel_size: 1 + ringattn_parallel_size: 1 + data_parallel_replicate_size: 1 + data_parallel_shard_size: 8 + expert_parallel_size: 8 + + num_train_epochs: 5 + max_steps: 3 + + 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: all_ranks + 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/gpt_oss_120b_ep8.yaml b/examples/local/dummy/configs/full/gpt_oss_120b_ep8.yaml new file mode 100644 index 00000000..caed0a8a --- /dev/null +++ b/examples/local/dummy/configs/full/gpt_oss_120b_ep8.yaml @@ -0,0 +1,54 @@ +model: + model_path: unsloth/gpt-oss-20b-BF16 + attn_implementation: eager + moe_implementation: native + +data: + datasets: + - path: dummy + type: tokenized + max_seq_len: 4000 + select_columns: [input_ids, labels] + dataset_prepared_path: last_prepared_dataset + sample_packing_method: sequential + sample_packing_sequence_len: 4000 + dataloader_num_workers: 4 + dataloader_prefetch_factor: 2 + dataloader_pin_memory: true + +train: + output_dir: outputs/gpt-oss-120b-ep8 + data_parallel_mode: fsdp2 + tensor_parallel_size: 1 + ulysses_parallel_size: 1 + ringattn_parallel_size: 1 + data_parallel_replicate_size: 1 + data_parallel_shard_size: 8 + expert_parallel_size: 8 + + num_train_epochs: 5 + + micro_batch_size: 1 + gradient_accumulation_steps: 1 + + optimizer: muon + muon_lr: 1.0e-4 + 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: all_ranks + 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/gpt_oss_20b_ep8.yaml b/examples/local/dummy/configs/full/gpt_oss_20b_ep8.yaml new file mode 100644 index 00000000..92dfa3fe --- /dev/null +++ b/examples/local/dummy/configs/full/gpt_oss_20b_ep8.yaml @@ -0,0 +1,55 @@ +model: + model_path: unsloth/gpt-oss-20b-BF16 + attn_implementation: eager + moe_implementation: native + +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/gpt-oss-20b-ep8 + data_parallel_mode: fsdp2 + tensor_parallel_size: 1 + ulysses_parallel_size: 1 + ringattn_parallel_size: 1 + data_parallel_replicate_size: 1 + data_parallel_shard_size: 8 + expert_parallel_size: 8 + + num_train_epochs: 5 + max_steps: 3 + + micro_batch_size: 1 + gradient_accumulation_steps: 1 + + optimizer: muon + muon_lr: 1.0e-4 + 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: all_ranks + 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.yaml b/examples/local/dummy/configs/full/llama3_8b.yaml index 76c7a5b6..ec7353ff 100644 --- a/examples/local/dummy/configs/full/llama3_8b.yaml +++ b/examples/local/dummy/configs/full/llama3_8b.yaml @@ -45,7 +45,7 @@ train: enable_full_shard: true enable_activation_offload: false init_device: meta - load_weights_mode: broadcast + load_weights_mode: grouped enable_full_determinism: false empty_cache_steps: 500 ckpt_manager: dcp diff --git a/examples/local/dummy/configs/full/llama3_8b_pp2.yaml b/examples/local/dummy/configs/full/llama3_8b_pp2.yaml index b4f6d869..743a45c5 100644 --- a/examples/local/dummy/configs/full/llama3_8b_pp2.yaml +++ b/examples/local/dummy/configs/full/llama3_8b_pp2.yaml @@ -43,7 +43,7 @@ train: enable_full_shard: true enable_activation_offload: false init_device: meta - load_weights_mode: broadcast + load_weights_mode: grouped enable_full_determinism: false empty_cache_steps: 500 ckpt_manager: dcp diff --git a/examples/local/dummy/configs/full/llama3_8b_tp4.yaml b/examples/local/dummy/configs/full/llama3_8b_tp4.yaml index e990cf43..0020abfd 100644 --- a/examples/local/dummy/configs/full/llama3_8b_tp4.yaml +++ b/examples/local/dummy/configs/full/llama3_8b_tp4.yaml @@ -44,7 +44,7 @@ train: enable_full_shard: true enable_activation_offload: false init_device: meta - load_weights_mode: broadcast + load_weights_mode: grouped enable_full_determinism: false empty_cache_steps: 500 ckpt_manager: dcp diff --git a/examples/local/dummy/configs/full/qwen3_30b_a3b_pp2_ep4_cp4_muon.yaml b/examples/local/dummy/configs/full/qwen3_30b_a3b_pp2_ep4_cp4_muon.yaml index 91d911e0..2920d10d 100644 --- a/examples/local/dummy/configs/full/qwen3_30b_a3b_pp2_ep4_cp4_muon.yaml +++ b/examples/local/dummy/configs/full/qwen3_30b_a3b_pp2_ep4_cp4_muon.yaml @@ -47,7 +47,7 @@ train: enable_full_shard: true enable_activation_offload: false init_device: meta - load_weights_mode: broadcast + load_weights_mode: grouped enable_full_determinism: false empty_cache_steps: 500 ckpt_manager: dcp diff --git a/examples/local/dummy/configs/full/qwen3_30b_a3b_pp2_ep4_muon.yaml b/examples/local/dummy/configs/full/qwen3_30b_a3b_pp2_ep4_muon.yaml index 5edb4416..f3c2aa42 100644 --- a/examples/local/dummy/configs/full/qwen3_30b_a3b_pp2_ep4_muon.yaml +++ b/examples/local/dummy/configs/full/qwen3_30b_a3b_pp2_ep4_muon.yaml @@ -46,7 +46,7 @@ train: enable_full_shard: true enable_activation_offload: false init_device: meta - load_weights_mode: broadcast + load_weights_mode: grouped enable_full_determinism: false empty_cache_steps: 500 ckpt_manager: dcp diff --git a/examples/local/dummy/configs/full/qwen3_32b.yaml b/examples/local/dummy/configs/full/qwen3_32b.yaml index aaf2f479..0c1e6f43 100644 --- a/examples/local/dummy/configs/full/qwen3_32b.yaml +++ b/examples/local/dummy/configs/full/qwen3_32b.yaml @@ -45,7 +45,7 @@ train: enable_full_shard: true enable_activation_offload: false init_device: meta - load_weights_mode: broadcast + load_weights_mode: grouped enable_full_determinism: false empty_cache_steps: 500 ckpt_manager: dcp diff --git a/examples/local/dummy/configs/full/qwen3_5_35b_a3b.yaml b/examples/local/dummy/configs/full/qwen3_5_35b_a3b.yaml index 71b8a9ed..8faf1d2c 100644 --- a/examples/local/dummy/configs/full/qwen3_5_35b_a3b.yaml +++ b/examples/local/dummy/configs/full/qwen3_5_35b_a3b.yaml @@ -46,7 +46,7 @@ train: enable_full_shard: true enable_activation_offload: false init_device: meta - load_weights_mode: broadcast + load_weights_mode: grouped enable_full_determinism: false ckpt_manager: dcp diff --git a/examples/local/dummy/configs/full/qwen3_8b.yaml b/examples/local/dummy/configs/full/qwen3_8b.yaml index 6c1bf34b..036a105d 100644 --- a/examples/local/dummy/configs/full/qwen3_8b.yaml +++ b/examples/local/dummy/configs/full/qwen3_8b.yaml @@ -46,7 +46,7 @@ train: enable_full_shard: true enable_activation_offload: false init_device: meta - load_weights_mode: broadcast + load_weights_mode: grouped enable_full_determinism: false empty_cache_steps: 500 ckpt_manager: dcp diff --git a/examples/local/dummy/configs/full/qwen3_8b_cp1_sp8.yaml b/examples/local/dummy/configs/full/qwen3_8b_cp1_sp8.yaml index 618111b3..2600265b 100644 --- a/examples/local/dummy/configs/full/qwen3_8b_cp1_sp8.yaml +++ b/examples/local/dummy/configs/full/qwen3_8b_cp1_sp8.yaml @@ -44,7 +44,7 @@ train: enable_full_shard: true enable_activation_offload: false init_device: meta - load_weights_mode: broadcast + load_weights_mode: grouped enable_full_determinism: false empty_cache_steps: 500 ckpt_manager: dcp diff --git a/examples/local/dummy/configs/full/qwen3_8b_cp2_sp4.yaml b/examples/local/dummy/configs/full/qwen3_8b_cp2_sp4.yaml index 0cecee24..36d463e0 100644 --- a/examples/local/dummy/configs/full/qwen3_8b_cp2_sp4.yaml +++ b/examples/local/dummy/configs/full/qwen3_8b_cp2_sp4.yaml @@ -44,7 +44,7 @@ train: enable_full_shard: true enable_activation_offload: false init_device: meta - load_weights_mode: broadcast + load_weights_mode: grouped enable_full_determinism: false empty_cache_steps: 500 ckpt_manager: dcp diff --git a/examples/local/dummy/configs/full/qwen3_8b_cp8_sp1.yaml b/examples/local/dummy/configs/full/qwen3_8b_cp8_sp1.yaml index 605ae471..d8528307 100644 --- a/examples/local/dummy/configs/full/qwen3_8b_cp8_sp1.yaml +++ b/examples/local/dummy/configs/full/qwen3_8b_cp8_sp1.yaml @@ -44,7 +44,7 @@ train: enable_full_shard: true enable_activation_offload: false init_device: meta - load_weights_mode: broadcast + load_weights_mode: grouped enable_full_determinism: false empty_cache_steps: 500 ckpt_manager: dcp diff --git a/examples/local/dummy/configs/full/qwen3_8b_muon.yaml b/examples/local/dummy/configs/full/qwen3_8b_muon.yaml index e8ebd65f..bdf8c7be 100644 --- a/examples/local/dummy/configs/full/qwen3_8b_muon.yaml +++ b/examples/local/dummy/configs/full/qwen3_8b_muon.yaml @@ -43,7 +43,7 @@ train: enable_full_shard: true enable_activation_offload: false init_device: meta - load_weights_mode: broadcast + load_weights_mode: grouped enable_full_determinism: false empty_cache_steps: 500 ckpt_manager: dcp diff --git a/examples/local/dummy/configs/full/qwen3_8b_pp2.yaml b/examples/local/dummy/configs/full/qwen3_8b_pp2.yaml index 010122db..acac27b0 100644 --- a/examples/local/dummy/configs/full/qwen3_8b_pp2.yaml +++ b/examples/local/dummy/configs/full/qwen3_8b_pp2.yaml @@ -44,7 +44,7 @@ train: enable_full_shard: true enable_activation_offload: false init_device: meta - load_weights_mode: broadcast + load_weights_mode: grouped enable_full_determinism: false empty_cache_steps: 500 ckpt_manager: dcp diff --git a/examples/local/dummy/configs/full/qwen3_8b_sp4_full_128k.yaml b/examples/local/dummy/configs/full/qwen3_8b_sp4_full_128k.yaml index b5b9f710..89f14bb2 100644 --- a/examples/local/dummy/configs/full/qwen3_8b_sp4_full_128k.yaml +++ b/examples/local/dummy/configs/full/qwen3_8b_sp4_full_128k.yaml @@ -42,7 +42,7 @@ train: enable_full_shard: true enable_activation_offload: false init_device: meta - load_weights_mode: broadcast + load_weights_mode: grouped enable_full_determinism: false empty_cache_steps: 500 ckpt_manager: dcp diff --git a/examples/local/dummy/configs/full/qwen3_8b_tp4_compile.yaml b/examples/local/dummy/configs/full/qwen3_8b_tp4_compile.yaml index 4f2c05ac..7c26aab0 100644 --- a/examples/local/dummy/configs/full/qwen3_8b_tp4_compile.yaml +++ b/examples/local/dummy/configs/full/qwen3_8b_tp4_compile.yaml @@ -45,7 +45,7 @@ train: enable_full_shard: true enable_activation_offload: false init_device: meta - load_weights_mode: broadcast + load_weights_mode: grouped enable_full_determinism: false empty_cache_steps: 500 ckpt_manager: dcp diff --git a/examples/local/dummy/configs/lora/llama3_8b_lora.yaml b/examples/local/dummy/configs/lora/llama3_8b_lora.yaml index b922bb59..bab67020 100644 --- a/examples/local/dummy/configs/lora/llama3_8b_lora.yaml +++ b/examples/local/dummy/configs/lora/llama3_8b_lora.yaml @@ -39,7 +39,7 @@ train: enable_gradient_checkpointing: true enable_full_shard: true init_device: meta - load_weights_mode: broadcast + load_weights_mode: grouped ckpt_manager: dcp save_steps: 0 save_hf_weights: false diff --git a/examples/local/dummy/configs/lora/qwen3_8b_lora.yaml b/examples/local/dummy/configs/lora/qwen3_8b_lora.yaml index 067f96d0..e68bf9cd 100644 --- a/examples/local/dummy/configs/lora/qwen3_8b_lora.yaml +++ b/examples/local/dummy/configs/lora/qwen3_8b_lora.yaml @@ -39,7 +39,7 @@ train: enable_gradient_checkpointing: true enable_full_shard: true init_device: meta - load_weights_mode: broadcast + load_weights_mode: grouped ckpt_manager: dcp save_steps: 0 save_hf_weights: false diff --git a/examples/local/dummy/configs/lora/qwen3_8b_lora_cp2_sp2.yaml b/examples/local/dummy/configs/lora/qwen3_8b_lora_cp2_sp2.yaml index d6ac3855..78fcaef5 100644 --- a/examples/local/dummy/configs/lora/qwen3_8b_lora_cp2_sp2.yaml +++ b/examples/local/dummy/configs/lora/qwen3_8b_lora_cp2_sp2.yaml @@ -40,7 +40,7 @@ train: enable_gradient_checkpointing: true enable_full_shard: true init_device: meta - load_weights_mode: broadcast + load_weights_mode: grouped ckpt_manager: dcp save_steps: 0 save_hf_weights: false diff --git a/examples/local/dummy/configs/lora/qwen3_8b_lora_cp4.yaml b/examples/local/dummy/configs/lora/qwen3_8b_lora_cp4.yaml index dcfed25f..ee73c3dd 100644 --- a/examples/local/dummy/configs/lora/qwen3_8b_lora_cp4.yaml +++ b/examples/local/dummy/configs/lora/qwen3_8b_lora_cp4.yaml @@ -40,7 +40,7 @@ train: enable_gradient_checkpointing: true enable_full_shard: true init_device: meta - load_weights_mode: broadcast + load_weights_mode: grouped ckpt_manager: dcp save_steps: 0 save_hf_weights: false diff --git a/examples/local/dummy/configs/qlora/llama3_8b_qlora_nvfp4.yaml b/examples/local/dummy/configs/qlora/llama3_8b_qlora_nvfp4.yaml index ffd3a529..9158a2a7 100644 --- a/examples/local/dummy/configs/qlora/llama3_8b_qlora_nvfp4.yaml +++ b/examples/local/dummy/configs/qlora/llama3_8b_qlora_nvfp4.yaml @@ -40,7 +40,7 @@ train: enable_gradient_checkpointing: true enable_full_shard: true init_device: meta - load_weights_mode: broadcast + load_weights_mode: grouped ckpt_manager: dcp save_steps: 0 save_hf_weights: false diff --git a/examples/local/dummy/configs/qlora/qwen3_32b_fp8_prequant_qlora.yaml b/examples/local/dummy/configs/qlora/qwen3_32b_fp8_prequant_qlora.yaml index 770a7f40..f75fe593 100644 --- a/examples/local/dummy/configs/qlora/qwen3_32b_fp8_prequant_qlora.yaml +++ b/examples/local/dummy/configs/qlora/qwen3_32b_fp8_prequant_qlora.yaml @@ -51,7 +51,7 @@ train: enable_gradient_checkpointing: true enable_full_shard: true init_device: meta - load_weights_mode: broadcast + load_weights_mode: grouped ckpt_manager: dcp save_steps: 0 save_hf_weights: false diff --git a/examples/local/dummy/configs/qlora/qwen3_32b_nvfp4_prequant_qlora.yaml b/examples/local/dummy/configs/qlora/qwen3_32b_nvfp4_prequant_qlora.yaml index 85407eac..0a36fb1e 100644 --- a/examples/local/dummy/configs/qlora/qwen3_32b_nvfp4_prequant_qlora.yaml +++ b/examples/local/dummy/configs/qlora/qwen3_32b_nvfp4_prequant_qlora.yaml @@ -39,7 +39,7 @@ train: enable_gradient_checkpointing: true enable_full_shard: true init_device: meta - load_weights_mode: broadcast + load_weights_mode: grouped ckpt_manager: dcp save_steps: 0 save_hf_weights: false diff --git a/examples/local/dummy/configs/qlora/qwen3_32b_qlora_nvfp4.yaml b/examples/local/dummy/configs/qlora/qwen3_32b_qlora_nvfp4.yaml index 53941787..81e1c6ce 100644 --- a/examples/local/dummy/configs/qlora/qwen3_32b_qlora_nvfp4.yaml +++ b/examples/local/dummy/configs/qlora/qwen3_32b_qlora_nvfp4.yaml @@ -39,7 +39,7 @@ train: enable_gradient_checkpointing: true enable_full_shard: true init_device: meta - load_weights_mode: broadcast + load_weights_mode: grouped ckpt_manager: dcp save_steps: 0 save_hf_weights: false diff --git a/examples/local/dummy/configs/qlora/qwen3_8b_qlora_block_fp8.yaml b/examples/local/dummy/configs/qlora/qwen3_8b_qlora_block_fp8.yaml index d00e9dc4..facd7ebf 100644 --- a/examples/local/dummy/configs/qlora/qwen3_8b_qlora_block_fp8.yaml +++ b/examples/local/dummy/configs/qlora/qwen3_8b_qlora_block_fp8.yaml @@ -51,7 +51,7 @@ train: enable_gradient_checkpointing: true enable_full_shard: true init_device: meta - load_weights_mode: broadcast + load_weights_mode: grouped ckpt_manager: dcp save_steps: 0 save_hf_weights: false diff --git a/examples/local/dummy/configs/qlora/qwen3_8b_qlora_nvfp4.yaml b/examples/local/dummy/configs/qlora/qwen3_8b_qlora_nvfp4.yaml index 9718a7fc..bc11489f 100644 --- a/examples/local/dummy/configs/qlora/qwen3_8b_qlora_nvfp4.yaml +++ b/examples/local/dummy/configs/qlora/qwen3_8b_qlora_nvfp4.yaml @@ -40,7 +40,7 @@ train: enable_gradient_checkpointing: true enable_full_shard: true init_device: meta - load_weights_mode: broadcast + load_weights_mode: grouped ckpt_manager: dcp save_steps: 0 save_hf_weights: false diff --git a/examples/local/dummy/configs/qlora/qwen3_8b_qlora_nvfp4_pp2.yaml b/examples/local/dummy/configs/qlora/qwen3_8b_qlora_nvfp4_pp2.yaml index 8f99e126..5d853521 100644 --- a/examples/local/dummy/configs/qlora/qwen3_8b_qlora_nvfp4_pp2.yaml +++ b/examples/local/dummy/configs/qlora/qwen3_8b_qlora_nvfp4_pp2.yaml @@ -41,7 +41,7 @@ train: enable_gradient_checkpointing: true enable_full_shard: true init_device: meta - load_weights_mode: broadcast + load_weights_mode: grouped ckpt_manager: dcp save_steps: 0 save_hf_weights: false diff --git a/examples/local/dummy/configs/qlora/qwen3_8b_qlora_nvfp4_requant.yaml b/examples/local/dummy/configs/qlora/qwen3_8b_qlora_nvfp4_requant.yaml index 9a5ee12a..e0bfcf9a 100644 --- a/examples/local/dummy/configs/qlora/qwen3_8b_qlora_nvfp4_requant.yaml +++ b/examples/local/dummy/configs/qlora/qwen3_8b_qlora_nvfp4_requant.yaml @@ -40,7 +40,7 @@ train: enable_gradient_checkpointing: true enable_full_shard: true init_device: meta - load_weights_mode: broadcast + load_weights_mode: grouped ckpt_manager: dcp save_steps: 0 save_hf_weights: false diff --git a/examples/local/dummy/configs/qlora/qwen3_8b_qlora_nvfp4_sp4.yaml b/examples/local/dummy/configs/qlora/qwen3_8b_qlora_nvfp4_sp4.yaml index e753564c..a4466bc7 100644 --- a/examples/local/dummy/configs/qlora/qwen3_8b_qlora_nvfp4_sp4.yaml +++ b/examples/local/dummy/configs/qlora/qwen3_8b_qlora_nvfp4_sp4.yaml @@ -40,7 +40,7 @@ train: enable_gradient_checkpointing: true enable_full_shard: true init_device: meta - load_weights_mode: broadcast + load_weights_mode: grouped ckpt_manager: dcp save_steps: 0 save_hf_weights: false diff --git a/examples/server/configs/full/gpt_oss_120b_full.yaml b/examples/server/configs/full/gpt_oss_120b_full.yaml new file mode 100644 index 00000000..2500342a --- /dev/null +++ b/examples/server/configs/full/gpt_oss_120b_full.yaml @@ -0,0 +1,45 @@ +# Server-side configuration for XORL Training Server (GPT-OSS 120B full weights, bf16) +# +# Full weight RL fine-tuning of the 120B MoE model. +# 8 GPUs: EP=8, FSDP shard=8. Muon optimizer (1 state) fits in GPU memory. +# Uses native backend (torch._grouped_mm) with alltoall EP dispatch. + +model_path: unsloth/gpt-oss-120b-BF16 +tokenizer_path: unsloth/gpt-oss-120b-BF16 +attn_implementation: eager +moe_implementation: native +ep_dispatch: alltoall +train_router: false + +data_parallel_mode: fsdp2 +expert_parallel_size: 8 +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/gpt-oss-120b-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: 32000 +enable_packing: true + +skip_initial_checkpoint: true + +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_30b_a3b_lora.yaml b/examples/server/configs/lora/qwen3_30b_a3b_lora.yaml new file mode 100644 index 00000000..5da454fa --- /dev/null +++ b/examples/server/configs/lora/qwen3_30b_a3b_lora.yaml @@ -0,0 +1,75 @@ +# Server-side configuration for XORL Training Server +# Qwen3-30B-A3B routed-expert multi-LoRA RL (bf16 base + LoRA rank 16) +# +# Smallest clean Qwen MoE (128 experts, 8 active/token, gated SwiGLU experts, +# full softmax attention). Minimal reference for RL training with multi-LoRA over +# routed experts. +# +# 1 node (8× H100): EP=8 (16 experts/rank), Ulysses SP=1, FSDP shard=1. +# This is the TRAINER config; pair it with a multiLoRA-enabled rollout engine. + +# ============================================================================ +# Model Configuration +# ============================================================================ +model_path: Qwen/Qwen3-30B-A3B +tokenizer_path: Qwen/Qwen3-30B-A3B +attn_implementation: flash_attention_3 # eager, sdpa, native, flash_attention_3, flash_attention_4 +moe_implementation: triton # eager, triton, native, quack + +# ============================================================================ +# Parallelism Configuration +# ============================================================================ +data_parallel_mode: fsdp2 # ddp, fsdp, fsdp2 +expert_parallel_size: 8 +ulysses_parallel_size: 1 +data_parallel_replicate_size: 1 +data_parallel_shard_size: 1 # REQUIRED for LoRA + EP: LoRA expert modules assert ep_fsdp_size == 1 + +# ============================================================================ +# Memory & Performance +# ============================================================================ +enable_mixed_precision: true +enable_gradient_checkpointing: true +enable_full_shard: true +enable_activation_offload: false +init_device: meta # cpu, meta, cuda + +# ============================================================================ +# Checkpointing +# ============================================================================ +output_dir: outputs/Qwen3-30B-A3B-server-lora-rl +load_checkpoint_path: "" +ckpt_manager: dcp # torch, dcp + +# ============================================================================ +# Logging +# ============================================================================ +log_level: INFO + +# ============================================================================ +# Worker Configuration +# ============================================================================ +worker_connection_timeout: 60.0 +worker_max_retries: 3 + +# ============================================================================ +# Data Processing Configuration +# ============================================================================ +sample_packing_sequence_len: 32768 +enable_packing: true + +# ============================================================================ +# LoRA Configuration +# ============================================================================ +enable_lora: true +lora_rank: 16 +lora_alpha: 16 +# Attention (q/k/v/o) + routed-expert (gate/up/down) projections. +# The last three target the MoE experts — this is what makes it routed-expert multiLoRA. +lora_target_modules: ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"] +# max_lora_rank defaults to lora_rank; raise it if fleet adapters use higher ranks. +# moe_hybrid_shared_lora: false # set true (+ lora_export_format: sglang_shared_outer) for the shared_outer layout + +# Skip initial checkpoint +skip_initial_checkpoint: true +ce_mode: compiled diff --git a/examples/server/no_robot_sft/README.md b/examples/server/no_robot_sft/README.md index 32cdf51f..f4ee0842 100644 --- a/examples/server/no_robot_sft/README.md +++ b/examples/server/no_robot_sft/README.md @@ -123,7 +123,7 @@ client = service.create_lora_training_client( # Training step result = client.forward_backward(datum_list, loss_fn="cross_entropy").result() -client.optim_step(xorl_client.AdamParams(learning_rate=1e-4)).result() +client.optim_step(learning_rate=1e-4).result() # Save checkpoint client.save_state("/path/to/checkpoint") @@ -142,7 +142,7 @@ client.save_state("/path/to/checkpoint") |----------|-------------| | `forward_backward(data)` | Forward + backward pass, accumulates gradients | | `forward(data)` | Forward-only pass (validation, no gradients) | -| `optim_step(params)` | Apply accumulated gradients with optimizer | +| `optim_step(learning_rate=...)` | Apply accumulated gradients with the session's configured optimizer | | `save_state(path)` | Save full training state (model + optimizer) | | `save_lora(path)` | Save LoRA adapter weights only | | `load_state(path)` | Load training state | diff --git a/examples/server/no_robot_sft/run_sft.py b/examples/server/no_robot_sft/run_sft.py index 0990a518..202c4737 100644 --- a/examples/server/no_robot_sft/run_sft.py +++ b/examples/server/no_robot_sft/run_sft.py @@ -30,7 +30,7 @@ class Config: learning_rate: float = 1e-4 max_length: int = 32768 train_on_what: renderers.TrainOnWhat = renderers.TrainOnWhat.ALL_ASSISTANT_MESSAGES - # we currently don't support modifying lora rank from the client side + # The server must be started with max_lora_rank >= this requested rank. lora_rank: int = 64 save_every: int = 20 # 0 = disabled @@ -129,8 +129,6 @@ def main(config: Config): # Linear learning rate schedule lr_mult = max(0.0, 1.0 - step / n_train_batches) current_lr = config.learning_rate * lr_mult - adam_params = xorl_client.AdamParams(learning_rate=current_lr, beta1=0.9, beta2=0.95, eps=1e-8) - # Get training batch and convert to datums online batch_start = batch_idx * config.batch_size batch_end = min((batch_idx + 1) * config.batch_size, len(train_dataset)) @@ -148,7 +146,7 @@ def main(config: Config): # Training step fwd_bwd_future = training_client.forward_backward(batch, loss_fn="cross_entropy") - optim_step_future = training_client.optim_step(adam_params) + optim_step_future = training_client.optim_step(learning_rate=current_lr) fwd_bwd_result = fwd_bwd_future.result() _optim_result = optim_step_future.result() diff --git a/examples/server/password_memorization/run_password_test.py b/examples/server/password_memorization/run_password_test.py index bfed180c..e4e11cd4 100644 --- a/examples/server/password_memorization/run_password_test.py +++ b/examples/server/password_memorization/run_password_test.py @@ -152,7 +152,9 @@ 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}, timeout=30) + resp = requests.post( + f"{train_url}/add_inference_endpoint", json={"host": host, "port": port, "worker_port": port}, timeout=30 + ) resp.raise_for_status() result = resp.json() si = result.get("endpoint", {}).get("server_info", {}) if result else {} @@ -173,7 +175,7 @@ def train_step(train_url, data, lr): loss = fb_result.get("metrics", {}).get("loss:mean", "N/A") opt = requests.post( f"{train_url}/api/v1/optim_step", - json={"model_id": MODEL_ID, "adam_params": {"learning_rate": lr}, "gradient_clip": 1.0}, + json={"model_id": MODEL_ID, "learning_rate": lr, "gradient_clip": 1.0}, timeout=30, ) opt.raise_for_status() diff --git a/examples/server/password_memorization/run_train_and_infer.py b/examples/server/password_memorization/run_train_and_infer.py index cbdde624..e7212a97 100644 --- a/examples/server/password_memorization/run_train_and_infer.py +++ b/examples/server/password_memorization/run_train_and_infer.py @@ -105,11 +105,9 @@ def main(config: Config): start_time = time.time() metrics = {} - adam_params = xorl_client.AdamParams(learning_rate=config.learning_rate, beta1=0.9, beta2=0.95, eps=1e-8) - # Training step fwd_bwd_future = training_client.forward_backward(datums, loss_fn="cross_entropy") - optim_step_future = training_client.optim_step(adam_params) + optim_step_future = training_client.optim_step(learning_rate=config.learning_rate) fwd_bwd_result = fwd_bwd_future.result() optim_result = optim_step_future.result() diff --git a/pyproject.sglang.toml b/pyproject.sglang.toml index f9cd32d3..1bebb45a 100644 --- a/pyproject.sglang.toml +++ b/pyproject.sglang.toml @@ -38,6 +38,8 @@ dependencies = [ "pyzmq", "msgpack", "fastapi", + # P2P / Mooncake weight sync + "mooncake-transfer-engine==0.3.9", "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", diff --git a/pyproject.toml b/pyproject.toml index 4cb41d74..4d134f14 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -18,11 +18,17 @@ dependencies = [ "packaging>=23.0,<26.0", "tenacity>=8.0.0", "torchdata>=0.8.0,<1.0", + "tiktoken", "transformers[torch]>=5.0", "psutil", "wandb", "safetensors", "einops", + # GatedDeltaNet kernels (replaces the vendored fork under ops/linear_attention/ops). + "flash-linear-attention @ git+https://github.com/fla-org/flash-linear-attention.git@97bcb883", + # FLA routes the gated-delta backward to TileLang on Hopper/Triton>=3.4. + "tilelang==0.1.11", + "apache-tvm-ffi==0.1.11", "numpy<2.4", "numba", "pydantic", @@ -30,6 +36,8 @@ dependencies = [ "pyzmq", "msgpack", "fastapi", + # P2P / Mooncake weight sync + "mooncake-transfer-engine==0.3.9", "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", @@ -37,7 +45,7 @@ dependencies = [ "triton @ https://download.pytorch.org/whl/triton-3.6.0-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%2B20260217.cu129torch2100cxx11abitrue.fec3a6-cp39-abi3-linux_x86_64.whl", - "flash-attn-cute @ git+https://github.com/Dao-AILab/flash-attention.git@5678dd909aca97f925957dab716022046fb1e44f#subdirectory=flash_attn/cute", + "flash-attn-4 @ git+https://github.com/Dao-AILab/flash-attention.git@59f01d6e1a1655a148ed4b22b5d4fbb9da2c2cf0#subdirectory=flash_attn/cute", ] @@ -114,6 +122,10 @@ markers = [ "dataloader: Tests for data loaders", ] testpaths = ["tests"] +# distributed_utils.py and other shared test helpers live alongside the +# distributed tests; put that directory on sys.path so test files can +# import them by bare module. +pythonpath = ["tests/distributed"] norecursedirs = ["submodules"] filterwarnings = [ "ignore::DeprecationWarning:multiprocessing.popen_fork", diff --git a/src/xorl/arguments.py b/src/xorl/arguments.py index 27c8a10d..42ba21b5 100644 --- a/src/xorl/arguments.py +++ b/src/xorl/arguments.py @@ -25,6 +25,7 @@ import torch import yaml +from .ops.loss import CrossEntropyMode from .utils import logging from .utils.checkpoint_utils import get_checkpoint_path @@ -46,7 +47,10 @@ def _detect_repo_commit() -> Optional[str]: ) commit = result.stdout.strip() return commit or None - except Exception: + except (subprocess.SubprocessError, FileNotFoundError, OSError): + # SubprocessError covers CalledProcessError when not a git repo; + # FileNotFoundError if `git` is not on PATH; OSError for other + # process-spawn failures. return None @@ -437,6 +441,10 @@ class ModelArguments: "'flash_attention_3': FA3 (Hopper). 'flash_attention_4': FA4 CUTE (Hopper+Blackwell)." }, ) + flash_attention_deterministic: bool = field( + default=False, + metadata={"help": "Request FlashAttention deterministic backward kernels when available."}, + ) moe_implementation: Optional[Literal[None, "eager", "triton", "native", "quack"]] = field( default=None, metadata={ @@ -454,6 +462,20 @@ class ModelArguments: "Disabled by default and must remain False when ep_dispatch='deepep'." }, ) + freeze_router: bool = field( + default=False, + metadata={"help": "Freeze MoE router weights during training."}, + ) + record_routing_weights: bool = field( + default=True, + metadata={ + "help": "Cache routing weights on the forward pass so they can override the " + "regathered weights during checkpoint recompute. Needed only when the " + "attention forward is non-deterministic across recompute (otherwise the " + "regather produces identical weights). Disabling skips the per-layer pinned " + "CPU allocation + D2H/H2D copies on every step." + }, + ) deepep_buffer_size_gb: float = field( default=2.0, metadata={"help": "DeepEP buffer size in GB (effective when ep_dispatch='deepep')."}, @@ -534,6 +556,17 @@ class TrainingArguments: default=0, metadata={"help": "L2 regularization strength."}, ) + cautious_weight_decay: bool = field( + default=False, + metadata={ + "help": "Apply Cautious Weight Decay (Chen et al., arXiv:2510.12402): " + "mask the decoupled decay term by I(u_t * x_t >= 0) so decay only acts " + "on coordinates whose update aligns with the parameter sign. " + "Supported with optimizer in {adamw, anyprecision_adamw, signsgd, muon}; " + "with adamw, routes to AnyPrecisionAdamW (fp32 state) since the fused " + "kernel has no per-coordinate decay hook." + }, + ) no_decay_modules: List[str] = field( default_factory=list, metadata={"help": "Modules without weight decay, for example, RMSNorm."}, @@ -543,10 +576,11 @@ class TrainingArguments: metadata={"help": "Parameters without weight decay, for example, bias."}, ) - optimizer: Literal["adamw", "anyprecision_adamw", "sgd", "signsgd", "muon"] = field( + optimizer: Literal["adamw", "anyprecision_adamw", "sgd", "signsgd", "distsignsgd", "muon"] = field( default="adamw", metadata={ - "help": "Optimizer type. 'signsgd' is a state-free sign update; 'muon' uses " + "help": "Optimizer type. 'signsgd' is a local state-free sign update; " + "'distsignsgd' signs gradients before FSDP2 reduction; 'muon' uses " "Newton-Schulz orthogonalization for 2D+ weight matrices." }, ) @@ -581,10 +615,12 @@ class TrainingArguments: }, ) muon_ns_algorithm: Literal["standard_newton_schulz", "gram_newton_schulz"] = field( - default="standard_newton_schulz", + default="gram_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." + "help": "Newton-Schulz backend for Muon. 'gram_newton_schulz' (default) batches across " + "MoE experts via baddbmm and is ~2x faster on Qwen3.5-style MoE; 'standard_newton_schulz' " + "uses the PyTorch upstream path (also batched in xorl) for bit-exact equivalence with " + "torch.optim._muon." }, ) muon_ns_use_quack_kernels: bool = field( @@ -608,6 +644,13 @@ class TrainingArguments: "A value of 2 means restart after the second iteration." }, ) + muon_grouped_gram_ns_fp32_byte_limit: int = field( + default=512 * 1024**2, + metadata={ + "help": "Maximum fp32 scratch bytes per grouped Muon Gram Newton-Schulz batch before chunking. " + "Lower values reduce peak optimizer scratch memory at the cost of more launches." + }, + ) muon_grad_dtype: Optional[Literal["fp32", "bf16"]] = field( default=None, metadata={ @@ -629,6 +672,17 @@ class TrainingArguments: "Intended for debugging and ablations." }, ) + muon_distributed_mode: Literal["shard_local", "full_gradient"] = field( + default="shard_local", + metadata={ + "help": "How Muon handles Newton-Schulz on FSDP2/EP-sharded DTensor params. " + "'shard_local': run NS on each rank's local shard (cheap, approximate). " + "'full_gradient': all-gather the post-momentum update to the full matrix, " + "run NS on the full matrix on every rank in the param's mesh, slice back to " + "the local shard. Recovers exact Muon at the cost of a per-step all-gather and " + "redundant NS compute. Implements the dense path of DeepSeek V4 §3.5.1." + }, + ) @property def optimizer_kwargs(self) -> Dict[str, Any]: @@ -643,6 +697,7 @@ def optimizer_kwargs(self) -> Dict[str, Any]: 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 + kwargs["muon_grouped_gram_ns_fp32_byte_limit"] = self.muon_grouped_gram_ns_fp32_byte_limit 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 @@ -658,12 +713,37 @@ def optimizer_kwargs(self) -> Dict[str, Any]: kwargs["muon_update_dtype"] = torch.float32 if self.muon_force_momentum_path: kwargs["muon_force_momentum_path"] = True + if self.muon_distributed_mode != "shard_local": + kwargs["muon_distributed_mode"] = self.muon_distributed_mode return kwargs max_grad_norm: float = field( default=1.0, metadata={"help": "Clip value for gradient norm."}, ) + softmax_auxiliary_loss: bool = field( + default=False, + metadata={ + "help": "Add a Z-loss auxiliary term on the LM-head logits: " + "mean(logsumexp(logits)^2) over valid tokens. Encourages the partition " + "function log(Z) to stay near zero, which stabilizes training at large " + "vocab and high learning rate (PaLM-style)." + }, + ) + auxiliary_loss_multiplier: float = field( + default=1e-5, + metadata={"help": "Coefficient for the softmax auxiliary (Z-)loss when softmax_auxiliary_loss is enabled."}, + ) + ce_mode: CrossEntropyMode = field( + default="compiled", + metadata={ + "help": "Cross-entropy computation mode for the local-trainer path. " + "'compiled' (default): torch.compile + auto_chunker, avoids materializing " + "the full [batch*seq, vocab] logits tensor. 'eager': F.cross_entropy that " + "materializes logits. The server path uses ServerArguments.ce_mode " + "(same semantics, separate flow)." + }, + ) micro_batch_size: int = field( default=1, metadata={"help": "Micro batch size. The number of samples per iteration on each device."}, @@ -753,10 +833,17 @@ def moe_recomputed(self) -> bool: "help": "Device to initialize model weights. 1. `cpu`: Init parameters on CPU in rank0 only. 2. `cuda`: Init parameters on GPU. 3. `meta`: Init parameters on meta. 4. `npu`: Init parameters on Ascend NPU." }, ) - load_weights_mode: Literal["broadcast", "all_ranks"] = field( - default="broadcast", + load_weights_mode: Literal["all_ranks", "grouped", "skip"] = field( + default="grouped", metadata={ - "help": "Weight loading mode. 'broadcast': rank0 reads weights and broadcasts to other ranks (default, avoids disk I/O bottleneck). 'all_ranks': every rank reads weights from disk independently." + "help": ( + "Weight loading mode. 'grouped' (default): one reader per node fan-outs dense " + "weights inside the node and one reader per EP-FSDP group fan-outs experts; when " + "grouped fanout groups are unavailable, it falls back to rank-0 loading. " + "'all_ranks': every rank reads weights from disk independently. 'skip': skip HF " + "weight loading entirely; use with load_checkpoint_path to materialize parameters " + "from a DCP checkpoint instead." + ) }, ) enable_full_determinism: bool = field( @@ -843,6 +930,17 @@ def moe_recomputed(self) -> bool: default="all", metadata={"help": "How to fold SP into FSDP: 'all' (ulysses+ring), 'ulysses_only', 'ring_only', or 'none'."}, ) + moe_grad_reduce_mode: Literal["reduce_scatter", "bf16_a2a_fp32_sum"] = field( + default="reduce_scatter", + metadata={ + "help": "Reduce-scatter strategy for MoE expert gradients on the ep_fsdp mesh dim. " + "'reduce_scatter': default NCCL reduce-scatter in FSDP's reduce_dtype (FP32). " + "'bf16_a2a_fp32_sum': stochastic-round FP32 grads to BF16, all-to-all across the " + "ep_fsdp group, then sum the received per-rank chunks locally in FP32. Halves " + "comm volume vs FP32 reduce-scatter while preserving FP32 accumulation precision. " + "Implements the MoE comm path of DeepSeek V4 §3.5.1. No effect on non-EP modules." + }, + ) ckpt_manager: Literal["dcp"] = field( default="dcp", metadata={"help": "Checkpoint manager."}, @@ -982,9 +1080,12 @@ 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}." - ) + if self.data_parallel_size != self.data_parallel_replicate_size * self.data_parallel_shard_size: + raise ValueError( + f"data_parallel_size ({self.data_parallel_size}) should equal " + f"data_parallel_replicate_size ({self.data_parallel_replicate_size}) * " + f"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: @@ -1016,12 +1117,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." - ) + if self.expert_parallel_size != 1 and self.init_device == "cpu": + raise ValueError( + "cpu init is not supported when expert parallelism is enabled. " + "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" + if self.data_parallel_mode == "fsdp2" and self.init_device != "meta": + raise ValueError("Please use init_device: meta for FSDP2 training") if self.load_checkpoint_path == "auto": self.load_checkpoint_path = get_checkpoint_path( @@ -1030,6 +1133,19 @@ def __post_init__(self): ckpt_manager=self.ckpt_manager, ) + if self.load_weights_mode not in {"grouped", "all_ranks", "skip"}: + raise ValueError( + f"Unsupported load_weights_mode={self.load_weights_mode!r}. Expected one of: grouped, all_ranks, skip." + ) + + if self.load_weights_mode == "skip" and not self.load_checkpoint_path: + raise ValueError( + "load_weights_mode='skip' skips HF weight loading and relies on " + "load_checkpoint_path to materialize parameters from a DCP checkpoint. " + "Set load_checkpoint_path (e.g. to the output of scripts/convert_checkpoint.py) " + "or choose a different load_weights_mode." + ) + # save paths self.save_checkpoint_path = os.path.join(self.output_dir, "checkpoints") self.step2token_path = os.path.join(self.output_dir, "step2token.json") @@ -1049,13 +1165,16 @@ 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!" - ) + if self.enable_full_determinism: + raise ValueError( + "allow_cuda_launch_blocking is disabled but enable_full_determinism is enabled. " + "enable_full_determinism would set CUDA_LAUNCH_BLOCKING=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!" - ) + if cuda_launch_blocking_val == "1": + raise ValueError( + "CUDA_LAUNCH_BLOCKING=1 is set in the environment but allow_cuda_launch_blocking is not enabled." + ) @dataclass @@ -1157,6 +1276,16 @@ class LoRAArguments: default=1.0, metadata={"help": "Scale factor for AQN noise magnitude."}, ) + moe_hybrid_shared_lora: bool = field( + default=False, + metadata={ + "help": "Route MoE LoRA injection through the hybrid-shared " + "(group-GEMM) path that shares lora_A across gate/up and " + "lora_B across down. Only meaningful for MoE models. " + "Read by Trainer._inject_lora; without this field on the " + "dataclass any local LoRA training run AttributeErrors out." + }, + ) @dataclass @@ -1279,8 +1408,8 @@ def parse_args(rootclass: T) -> T: base_to_subclass[base] = subclass.default_factory try: type_hints: Dict[str, type] = get_type_hints(subclass.default_factory) - except Exception: - raise RuntimeError(f"Type resolution failed for {subclass.default_factory}.") + except (NameError, TypeError, AttributeError) as e: + raise RuntimeError(f"Type resolution failed for {subclass.default_factory}.") from e for attr in fields(subclass.default_factory): if not attr.init: @@ -1316,8 +1445,8 @@ def parse_args(rootclass: T) -> T: base = subclass.name try: type_hints: Dict[str, type] = get_type_hints(subclass.default_factory) - except Exception: - raise RuntimeError(f"Type resolution failed for {subclass.default_factory}.") + except (NameError, TypeError, AttributeError) as e: + raise RuntimeError(f"Type resolution failed for {subclass.default_factory}.") from e for attr in fields(subclass.default_factory): if not attr.init: diff --git a/src/xorl/checkpoint/checkpointer.py b/src/xorl/checkpoint/checkpointer.py index c91646c6..7c5670ad 100644 --- a/src/xorl/checkpoint/checkpointer.py +++ b/src/xorl/checkpoint/checkpointer.py @@ -2,6 +2,7 @@ import json import os from abc import ABC, abstractmethod +from collections import OrderedDict from typing import Any, Dict, List, Optional, Set import torch @@ -215,6 +216,46 @@ def state_dict(self): return model_state_dict + @torch.no_grad() + def reference_state_dict(self): + """Collect a lightweight state dict of live params/buffers without DCP materialization. + + This is intended for direct safetensors export paths where we only need + references to the current model tensors and will materialize them one at + a time during save. + """ + model_state_dict: "OrderedDict[str, torch.Tensor]" = OrderedDict() + modules = dict(self.model.named_modules(remove_duplicate=False)) + + for name, parameter in self.model.named_parameters(remove_duplicate=False): + if parameter is not None: + model_state_dict[name] = parameter + + for name, buffer in self.model.named_buffers(remove_duplicate=False): + if buffer is None: + continue + module_name, _, buffer_name = name.rpartition(".") + parent_module = modules[module_name] if module_name else self.model + if buffer_name in getattr(parent_module, "_non_persistent_buffers_set", set()): + continue + model_state_dict[name] = buffer + + if self.should_ep_aware: + logger.info_rank0( + "Collecting lightweight model tensor references from ModelState wrapper, " + "restoring EP dim for Experts module" + ) + model_state_dict = self.get_state_dict_with_ep_dim(model_state_dict) + + if self.exclude_keys: + model_state_dict = OrderedDict((k, v) for k, v in model_state_dict.items() if k not in self.exclude_keys) + + if self.save_lora_only: + model_state_dict = OrderedDict((k, v) for k, v in model_state_dict.items() if "lora_" in k) + logger.info_rank0(f"LoRA-only save: keeping {len(model_state_dict)} LoRA parameters") + + return model_state_dict + @torch.no_grad() def load_state_dict(self, state_dict): """ @@ -595,8 +636,19 @@ def load( logger.info_rank0(f"LoRA-only checkpoint: excluding {len(exclude_keys)} non-LoRA keys from load") load_state = {"model": ModelState(state["model"], exclude_keys=exclude_keys)} + has_optimizer_state = False if "optimizer" in state and state["optimizer"] is not None: + try: + dcp_metadata = FileSystemReader(checkpoint_dir).read_metadata() + has_optimizer_state = any(key.startswith("optimizer") for key in dcp_metadata.state_dict_metadata) + except Exception as exc: + logger.warning_rank0(f"Could not inspect DCP optimizer metadata at {checkpoint_dir}: {exc}") + has_optimizer_state = True + + if has_optimizer_state: load_state["optimizer"] = OptimizerState(model=state["model"], optimizer=state["optimizer"]) # type: ignore[index] + elif "optimizer" in state and state["optimizer"] is not None: + logger.info_rank0(f"No optimizer state found in {checkpoint_dir}; loading model state only.") dcp.load( state_dict=load_state, @@ -607,9 +659,11 @@ def load( if "extra_state" in state: extra_state_dir = os.path.join(checkpoint_dir, _EXTRA_STATE_DIR) - os.makedirs(extra_state_dir, exist_ok=True) extra_state_path = os.path.join(extra_state_dir, _EXTRA_STATE_FORMAT.format(dist.get_rank())) - state["extra_state"] = torch.load(extra_state_path, weights_only=False) + if os.path.exists(extra_state_path): + state["extra_state"] = torch.load(extra_state_path, weights_only=False) + else: + logger.info_rank0(f"No extra_state found at {extra_state_path}, starting fresh.") logger.info_rank0(f"Loaded checkpoint from {checkpoint_dir}") diff --git a/src/xorl/cli/direct_train.py b/src/xorl/cli/direct_train.py deleted file mode 100644 index e6bf5606..00000000 --- a/src/xorl/cli/direct_train.py +++ /dev/null @@ -1,844 +0,0 @@ -# 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 -from typing import Any, Dict, List, Optional - -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 -from xorl.data.prepare.prepare_datasets import prepare_datasets -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.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.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, - maybe_merge_lora, - 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__) -_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(): - args = parse_args(Arguments) - 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, - dp_shard_size=args.train.data_parallel_shard_size, - tp_size=args.train.tensor_parallel_size, - ep_size=args.train.expert_parallel_size, - pp_size=args.train.pipeline_parallel_size, - ulysses_size=args.train.ulysses_parallel_size, - ringattn_size=args.train.ringattn_parallel_size, - dp_mode=args.train.data_parallel_mode, - cp_fsdp_mode=args.train.cp_fsdp_mode, - ) - - # Initialize DTensor RNG tracker with run_state_sync=False to prevent a - # world-group broadcast that deadlocks when PP stages run asynchronously. - # Same approach as torchtitan (see torchtitan/distributed/utils.py:set_determinism). - ps = get_parallel_state() - if ps.device_mesh is not None: - torch.distributed.tensor._random.manual_seed(args.train.seed, ps.device_mesh) - - logger.info_rank0("Prepare data") - tokenizer = build_tokenizer(args.model.tokenizer_path) - - # Load the datasets - train_dataset, eval_dataset = prepare_datasets(args, tokenizer) - - train_dataloader = DataLoaderBuilder( - dataset=train_dataset, - micro_batch_size=args.train.micro_batch_size, - gradient_accumulation_steps=args.train.gradient_accumulation_steps, - num_workers=args.data.dataloader_num_workers, - drop_last=args.data.dataloader_drop_last, - pin_memory=args.data.dataloader_pin_memory, - prefetch_factor=args.data.dataloader_prefetch_factor, - 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 - if args.train.max_steps is not None: - total_train_steps = min(total_train_steps, args.train.max_steps) - logger.info_rank0(f"Train steps per epoch: {train_steps_per_epoch}, Total train steps: {total_train_steps}") - - # Convert save_epochs (fractional) to a step interval - save_epoch_steps = int(args.train.save_epochs * train_steps_per_epoch) if args.train.save_epochs else 0 - 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=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 - helper.print_device_mem_info("VRAM usage after building model") - - # Unfuse QKV projections if merge_qkv=False so each projection is handled independently. - if not args.model.merge_qkv: - for layer in model.model.layers: - if hasattr(layer, "self_attn") and hasattr(layer.self_attn, "unfuse_for_tp"): - layer.self_attn.unfuse_for_tp() - logger.info_rank0("Unfused QKV projections (merge_qkv=False)") - - # QLoRA injection: replace target nn.Linear with QLoRALinear. - # With meta init, quantization is deferred — weights stay as nn.Parameter - # so FSDP can load them normally. After FSDP loading, maybe_quantize_qlora() - # converts them to uint8 buffers. - # For pre-quantized checkpoints (NVFP4 modelopt format), quantization is - # skipped entirely — packed weights + scales are loaded directly. - is_prequantized = False - checkpoint_quant_format = None - exclude_modules = set() - if args.lora.enable_qlora: - if detect_prequantized_nvfp4(args.model.model_path): - is_prequantized = True - checkpoint_quant_format = "nvfp4" - logger.info_rank0("Detected pre-quantized NVFP4 checkpoint") - elif detect_prequantized_block_fp8(args.model.model_path): - is_prequantized = True - checkpoint_quant_format = "block_fp8" - 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}") - elif is_prequantized: - exclude_modules = get_prequantized_exclude_modules(args.model.model_path) - if exclude_modules: - logger.info_rank0( - 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( - f"Cross-format conversion: checkpoint={checkpoint_quant_format}, " - f"target={args.lora.quant_format} — will dequantize and re-quantize" - ) - - inject_qlora_into_model( - model, - r=args.lora.lora_rank, - lora_alpha=args.lora.lora_alpha, - quant_format=args.lora.quant_format, - quant_group_size=args.lora.quant_group_size, - target_modules=args.lora.lora_target_modules, - checkpoint_quant_format=checkpoint_quant_format, - merge_qkv=args.model.merge_qkv, - exclude_modules=exclude_modules, - enable_aqn=args.lora.enable_aqn, - aqn_alpha=args.lora.aqn_alpha, - ) - # Store exclude_modules on model so checkpoint handler can use the - # same set (user-specified or auto-detected) instead of re-detecting. - if exclude_modules: - model._qlora_exclude_modules = exclude_modules - helper.print_device_mem_info("VRAM usage after QLoRA injection") - elif args.lora.enable_lora: - 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, - init_device=args.train.init_device, - weights_path=args.model.model_path, - enable_full_shard=args.train.enable_full_shard, - enable_mixed_precision=args.train.enable_mixed_precision, - enable_gradient_checkpointing=args.train.enable_gradient_checkpointing, - enable_compile=args.train.enable_compile, - basic_modules=model._no_split_modules + args.model.basic_modules, - enable_reentrant=args.train.enable_reentrant, - 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=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 - pp_enabled = isinstance(build_result, dict) - pp_stages = None - model_parts = None - has_first_stage = False - has_last_stage = False - if pp_enabled: - pp_stages = build_result["stages"] - model_parts = build_result["model_parts"] - has_first_stage = build_result["has_first_stage"] - has_last_stage = build_result["has_last_stage"] - model = model_parts[0] # primary model for optimizer etc. - else: - model = build_result - - # Deferred QLoRA quantization: now that FSDP has loaded weights into the - # weight parameters, quantize them into uint8 buffers and free the originals. - # For pre-quantized checkpoints, skip quantization — load packed weights directly. - if args.lora.enable_qlora: - if is_prequantized: - 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: - logger.info("Starting maybe_quantize_qlora...") - helper.print_device_mem_info("VRAM before QLoRA quantization") - maybe_quantize_qlora(model) - logger.info("Done maybe_quantize_qlora, starting MoE weight loading...") - helper.print_device_mem_info("VRAM after QLoRA linear quantization") - # Load and quantize MoE expert weights directly from checkpoint - # (bypasses FSDP to avoid OOM for large MoE models) - maybe_load_and_quantize_moe_qlora(model, args.model.model_path) - 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",), - ) - torch.cuda.empty_cache() - if removed > 0: - logger.info(f"Deregistered {removed} packed_weight_f32 params from FSDP2") - # Freeze all non-LoRA parameters (embeddings, norms, lm_head, etc.) - for name, param in model.named_parameters(): - if "lora_A" not in name and "lora_B" not in name: - param.requires_grad = False - helper.print_device_mem_info("VRAM usage after QLoRA quantization") - - optimizer = build_optimizer( - model, - lr=args.train.lr, - weight_decay=args.train.weight_decay, - fused=True, - optimizer_type=args.train.optimizer, - optimizer_dtype=args.train.optimizer_dtype, - optimizer_kwargs=args.train.optimizer_kwargs, - ) - if get_optimizer_pre_hook is not None: - optimizer_pre_hook = get_optimizer_pre_hook(model, model_config, args.train.data_parallel_mode) - optimizer.register_step_pre_hook(optimizer_pre_hook) - - lr_scheduler = build_lr_scheduler( - optimizer, - train_steps=total_train_steps, - lr=args.train.lr, - lr_min=args.train.lr_min, - lr_decay_style=args.train.lr_decay_style, - lr_decay_ratio=args.train.lr_decay_ratio, - lr_warmup_ratio=args.train.lr_warmup_ratio, - lr_start=args.train.lr_start, - ) - - if args.train.global_rank == 0: - # save model_assets before training - model_assets = [model_config, tokenizer] - save_model_assets(args.train.model_assets_dir, model_assets) - - if args.train.profile_this_rank: - profiler = helper.create_profiler( - start_step=args.train.profile_start_step, - end_step=args.train.profile_end_step, - trace_dir=args.train.profile_trace_dir, - record_shapes=args.train.profile_record_shapes, - profile_memory=args.train.profile_profile_memory, - with_stack=args.train.profile_with_stack, - global_rank=args.train.global_rank, - ) - profiler.start() - - start_epoch, start_step, global_step = 0, 0, 0 - save_checkpoint_path = None - environ_meter = helper.EnvironMeter( - config=model_config, - global_batch_size=args.train.global_batch_size, - empty_cache_steps=args.train.empty_cache_steps, - 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, - ) - - if args.train.load_checkpoint_path: - state = {"model": model, "optimizer": optimizer, "extra_state": {}} # cannot be None - Checkpointer.load(args.train.load_checkpoint_path, state) - global_step = state["extra_state"]["global_step"] - start_epoch = global_step // train_steps_per_epoch - start_step = global_step % train_steps_per_epoch - lr_scheduler.load_state_dict(state["extra_state"]["lr_scheduler"]) - train_dataloader.load_state_dict(state["extra_state"]["train_dataloader"]) - environ_meter.load_state_dict(state["extra_state"]["environ_meter"]) - torch.set_rng_state(state["extra_state"]["torch_rng_state"]) - if start_step == 0: # resume at the end of epoch - iter(train_dataloader) # clear resume state and prefetch data - - dist.barrier() - logger.info_rank0(f"Load distributed checkpoint from {args.train.load_checkpoint_path} successfully!") - - # Build PP schedule if pipeline parallelism is enabled - pp_schedule = None - pp_context = {} # mutable container for per-step state used by pp_loss_fn - if pp_enabled: - - @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 - ) - - def pp_loss_fn(pred, labels): - return _pp_ce_loss(pred, labels, pp_context["global_valid_tokens"]) - - pp_schedule = build_pipeline_schedule( - stages=pp_stages, - n_microbatches=args.train.gradient_accumulation_steps, - loss_fn=pp_loss_fn, - schedule_name=args.train.pipeline_parallel_schedule, - ) - logger.info_rank0(f"PP schedule built: {args.train.pipeline_parallel_schedule}") - - helper.empty_cache() - model_fwd_context, model_bwd_context = build_activation_offloading_context( - args.train.enable_activation_offload, args.train.enable_gradient_checkpointing, args.train.activation_gpu_limit - ) - model.train() - logger.info( - f"rank{args.train.local_rank} Start training, train_steps_per_epoch: {train_steps_per_epoch}, " - f"total_train_steps: {total_train_steps}, epochs: {args.train.num_train_epochs}" - ) - for epoch in range(start_epoch, args.train.num_train_epochs): - if hasattr(train_dataloader, "set_epoch"): - train_dataloader.set_epoch(epoch) - - # Compute actual steps this epoch, capped by max_steps - steps_this_epoch = train_steps_per_epoch - start_step - if args.train.max_steps is not None: - steps_this_epoch = min(steps_this_epoch, args.train.max_steps - global_step) - if steps_this_epoch <= 0: - break - - data_loader_tqdm = trange( - steps_this_epoch, - desc=f"Epoch {epoch + 1}/{args.train.num_train_epochs}", - total=start_step + steps_this_epoch, - initial=start_step, - disable=args.train.local_rank != 0, - ) - data_iterator = iter(train_dataloader) - for _ in range(start_step, train_steps_per_epoch): - if args.train.max_steps is not None and global_step >= args.train.max_steps: - logger.info_rank0(f"Reached max_steps={args.train.max_steps}, stopping training.") - break - global_step += 1 - - try: - micro_batches: List[Dict[str, Any]] = next(data_iterator) - except StopIteration: - logger.info(f"epoch:{epoch} Dataloader finished with drop_last {args.data.drop_last}") - break - - # Synchronize padding across DP ranks to prevent load imbalance - # Use fsdp_group (not world) when PP enabled to avoid NCCL conflicts - sync_group = ps.fsdp_group if pp_enabled else None - synchronize_micro_batch_padding(micro_batches, group=sync_group) - - if global_step == 1: - helper.print_example(example=micro_batches[0], rank=args.train.local_rank) - - total_loss = 0 - synchronize() - start_time = time.time() - - # compute global valid tokens across all ranks - global_valid_tokens = count_valid_tokens( - micro_batches, - group=ps.fsdp_group if pp_enabled else None, - ) - - optimizer.zero_grad() - - # AQN: pre-generate noise for all QLoRA layers on side streams. - # Runs async — overlaps with data prep below, so forward only - # pays the cheap addcmul cost per layer. - if args.lora.enable_aqn: - 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. - - use_routing_replay = ps.ep_size > 1 and args.train.moe_recomputed - - if pp_enabled: - # === Pipeline Parallel training path === - # Set global_valid_tokens for pp_loss_fn normalization - pp_context["global_valid_tokens"] = global_valid_tokens - - for micro_batch in micro_batches: - environ_meter.add(micro_batch) - - # 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) - - # 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. - - _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: - md = {} - if "position_ids" in mb: - md["position_ids"] = mb["position_ids"] - for key in _PP_FA_KEYS: - if key in mb: - md[key] = mb[key] - pp_metadata_list.append(md) - for model_part in model_parts: - model_part._pp_batch_metadata = deque(pp_metadata_list) - - # Only last stage computes loss - if has_last_stage: - targets = labels - losses = [] - else: - targets = None - losses = None - - # Routing replay: global stage = "replay_backward" so checkpoint - # recompute uses recorded routing. _pp_forward temporarily - # switches to "record" during each forward call. - if use_routing_replay: - set_replay_stage("replay_backward") - - # Run PP schedule (handles fwd/bwd for all microbatches) - if has_first_stage: - pp_schedule.step(input_ids, target=targets, losses=losses) - else: - pp_schedule.step(target=targets, losses=losses) - - if use_routing_replay: - set_replay_stage(None) - RoutingReplay.clear_all() - - # Compute loss for logging (losses already normalized by global_valid_tokens) - if has_last_stage: - total_loss = torch.sum(torch.stack(losses)).item() - loss_tensor = torch.tensor([total_loss], device=device) - else: - loss_tensor = torch.tensor([-1.0], device=device) - - # Share loss across PP stages (MAX broadcasts from last stage) - dist.all_reduce(loss_tensor, op=dist.ReduceOp.MAX, group=ps.pp_group) - total_loss = loss_tensor.item() - - del input_ids, labels - else: - # === Standard gradient accumulation path === - # Global stage for recompute; switched to "record" around forward. - if use_routing_replay: - set_replay_stage("replay_backward") - - for micro_batch in micro_batches: - environ_meter.add(micro_batch) - - micro_batch = { - k: v.to(get_device_type(), non_blocking=True) if isinstance(v, torch.Tensor) else v - for k, v in micro_batch.items() - } - - # Pop labels before forward (model no longer takes labels) - labels = micro_batch.pop("labels", None) - - if use_routing_replay: - set_replay_stage("record") - with model_fwd_context: - outputs = model(**micro_batch, use_cache=False, output_hidden_states=False) - - # 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, - ) - 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 = 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) - - local_valid_tokens = (labels != IGNORE_INDEX).sum() - ga_loss, _ = gradient_accumulate_loss(loss, local_valid_tokens, global_valid_tokens) - if use_routing_replay: - set_replay_stage("replay_backward") - - with model_bwd_context: - ga_loss.backward() - - # NOTE: Do NOT reset backward indices here — the backward_index - # must increment across micro-batches (entry 0 = MB0, entry 1 = MB1, etc.) - # reset_all_backward() would cause MB1's recompute to replay MB0's routing. - - loss_item = ga_loss.item() - total_loss += loss_item - - # Clean up tensors to free memory - del micro_batch, labels, loss, outputs, ga_loss - - if use_routing_replay: - set_replay_stage(None) - RoutingReplay.clear_all() - - # Sync gradients across ring/Ulysses dims not folded into FSDP - sync_sp_gradients(model, ps.sp_grad_sync_group) - - # Gradient clipping - grad_norm = clip_gradients( - model, - args.train.max_grad_norm, - pp_enabled=pp_enabled, - pp_group=ps.pp_group if pp_enabled else None, - ) - - optimizer.step() - lr_scheduler.step() - optimizer.zero_grad() - - # Periodic LoRA merge: absorb LoRA delta into base weights - maybe_merge_lora( - model, - enable_lora=args.lora.enable_lora, - enable_qlora=args.lora.enable_qlora, - merge_interval=args.lora.merge_lora_interval, - global_step=global_step, - optimizer=optimizer, - reset_optimizer=args.lora.reset_optimizer_on_merge, - ) - if hasattr(grad_norm, "full_tensor"): - grad_norm = grad_norm.full_tensor().item() - - # Collect mean loss and grad_norm across data parallel group for logging. - # For PP: total_loss is this rank's sum(losses)/global_valid_tokens where - # global_valid_tokens already includes all DP ranks. SUM across DP gives - # the correct global per-token loss. grad_norm is already consistent - # across all ranks (FSDP all-reduce + PP MAX), so just take mean (no-op). - if pp_enabled: - total_loss = all_reduce(total_loss, op="sum", group=ps.fsdp_group) - grad_norm = all_reduce(grad_norm, op="mean", group=ps.fsdp_group) - else: - total_loss, grad_norm = all_reduce((total_loss, grad_norm), group=ps.fsdp_group) - synchronize() - delta_time = time.time() - start_time - lr = max(lr_scheduler.get_last_lr()) - train_metrics = environ_meter.step(delta_time, global_step=global_step) - - tokens_per_sec = train_metrics.get("efficiency/tokens_per_second(K)", 0) * 1e3 - data_loader_tqdm.set_postfix_str( - f"loss={total_loss:.2f} gn={grad_norm:.2f} lr={lr:.1e} tok/s={tokens_per_sec:.0f}" - ) - data_loader_tqdm.update() - - if args.train.global_rank == 0: - if args.train.use_wandb and global_step % args.train.wandb_log_interval == 0: - 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: - profiler.step() - 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 - ) - if should_save: - helper.empty_cache() - save_checkpoint_path = os.path.join(args.train.save_checkpoint_path, f"global_step_{global_step}") - state = { - "model": model, - "optimizer": optimizer, - "extra_state": { - "global_step": global_step, - "lr_scheduler": lr_scheduler.state_dict(), - "train_dataloader": train_dataloader.state_dict(), - "environ_meter": environ_meter.state_dict(), - "torch_rng_state": torch.get_rng_state(), - }, - } - # Determine if we should save only LoRA params (base weights unchanged) - 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, - ) - - dist.barrier() - logger.info_rank0(f"Distributed checkpoint saved at {save_checkpoint_path} successfully!") - - data_loader_tqdm.close() - start_step = 0 - helper.print_device_mem_info(f"VRAM usage after epoch {epoch + 1}") - - synchronize() - - # Gather full model state via NCCL for HF save (all ranks must participate). - # This is much faster than the DCP round-trip (write to disk → read back) - # because NCCL AllGather is ~10-50 GB/s vs ~0.65 GB/s NFS. - is_lora_training = args.lora.enable_lora or args.lora.enable_qlora - save_peft_adapter = is_lora_training and args.lora.merge_lora_interval == 0 - - hf_model_state_dict = None - if args.train.save_hf_weights and not save_peft_adapter: - 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)) - - # release memory - del optimizer, lr_scheduler - helper.empty_cache() - - # save model in huggingface's format (rank 0 only) - 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_{global_step}", "hf_ckpt") - if save_peft_adapter: - # Save PEFT adapter format (LoRA-only, base weights unchanged) - - save_lora_checkpoint( - 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 - ) - del hf_model_state_dict - logger.info_rank0(f"Huggingface checkpoint saved at {hf_weights_path} successfully!") - - dist.barrier() - dist.destroy_process_group() - - -if __name__ == "__main__": - main() diff --git a/src/xorl/cli/preprocess.py b/src/xorl/cli/preprocess.py index e7ce52b3..ba5d3f1d 100644 --- a/src/xorl/cli/preprocess.py +++ b/src/xorl/cli/preprocess.py @@ -64,12 +64,12 @@ def load_config_without_validation(config_path: str) -> Arguments: def main(): """Preprocess datasets and save them to disk for later use in training.""" if len(sys.argv) < 2: - print("Usage: python -m xorl.cli.preprocess ") + sys.stderr.write("Usage: python -m xorl.cli.preprocess \n") sys.exit(1) config_path = sys.argv[1] if not os.path.exists(config_path): - print(f"Error: Config file not found: {config_path}") + sys.stderr.write(f"Error: Config file not found: {config_path}\n") sys.exit(1) # Load config without triggering distributed validation diff --git a/src/xorl/data/collators/packing_concat_collator.py b/src/xorl/data/collators/packing_concat_collator.py index 7c54041c..bd4666dc 100644 --- a/src/xorl/data/collators/packing_concat_collator.py +++ b/src/xorl/data/collators/packing_concat_collator.py @@ -1,4 +1,3 @@ -import logging from dataclasses import dataclass from typing import Dict, Sequence, Tuple @@ -10,9 +9,6 @@ from .base_collator import DataCollator -logger = logging.getLogger(__name__) - - def add_flash_attention_kwargs_from_position_ids( batch: Dict[str, "torch.Tensor"], ) -> Tuple["torch.Tensor", "torch.Tensor", "torch.Tensor", "torch.Tensor"]: @@ -73,8 +69,6 @@ def __call__(self, features: Sequence[Dict[str, "torch.Tensor"]]) -> Dict[str, " if not features: raise ValueError("PackingConcatCollator received empty features list") - logger = logging.getLogger(__name__) - # Input should be a flat list of dicts from FlattenCollator assert isinstance(features[0], dict), ( f"Expected dict from FlattenCollator, but got {type(features[0]).__name__}" @@ -82,7 +76,8 @@ def __call__(self, features: Sequence[Dict[str, "torch.Tensor"]]) -> Dict[str, " batch = {} for input_name in features[0].keys(): - # Handle 1D tensors (input_ids, labels, etc.) and 2D tensors (hidden_states, hidden_states_scale) + # Handle 1D tensors (input_ids, labels, etc.) and 2D tensors + # (hidden_states, teacher_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", @@ -92,6 +87,10 @@ def __call__(self, features: Sequence[Dict[str, "torch.Tensor"]]) -> Dict[str, " "target_tokens", "logprobs", "advantages", + "rollout_logprobs", + "teacher_ids", + "teacher_cache_indices", + "teacher_weights", ): # 1D tensors: concatenate along sequence dimension tensors = [feature[input_name] for feature in features] @@ -108,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"): + elif input_name in ("hidden_states", "teacher_hidden_states", "hidden_states_scale"): # 2D tensors: concatenate along sequence dimension (dim 0) tensors = [feature[input_name] for feature in features] @@ -172,15 +171,23 @@ def __call__(self, features: Sequence[Dict[str, "torch.Tensor"]]) -> Dict[str, " ) 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( - batch["hidden_states"], (0, 0, 0, pad_length), value=0.0 - ) - if "hidden_states_scale" in batch: - batch["hidden_states_scale"] = torch.nn.functional.pad( - batch["hidden_states_scale"], (0, 0, 0, pad_length), value=0.0 - ) + for key in ( + "target_tokens", + "logprobs", + "advantages", + "rollout_logprobs", + "teacher_ids", + "teacher_cache_indices", + "teacher_weights", + ): + if key in batch and batch[key].shape[-1] == seq_len: + pad_value = -100 if key == "target_tokens" else 0 + batch[key] = torch.nn.functional.pad(batch[key], (0, pad_length), value=pad_value) + + # Pad 2D sequence-aligned tensors. + for key in ("hidden_states", "teacher_hidden_states", "hidden_states_scale"): + if key in batch: + batch[key] = torch.nn.functional.pad(batch[key], (0, 0, 0, pad_length), value=0.0) # 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: diff --git a/src/xorl/data/collators/sequence_shard_collator.py b/src/xorl/data/collators/sequence_shard_collator.py index 6493b059..2c427b91 100644 --- a/src/xorl/data/collators/sequence_shard_collator.py +++ b/src/xorl/data/collators/sequence_shard_collator.py @@ -239,34 +239,54 @@ def __call__(self, batch: Dict[str, "torch.Tensor"]) -> Dict[str, "torch.Tensor" batch["labels"] = self.sp_slice(labels, dim=-1) batch["position_ids"] = position_ids # Keep full, not sliced - # Handle RL fields (target_tokens, logprobs, advantages) for importance_sampling - # These need to be padded and sliced the same way as labels - rl_fields = ["target_tokens", "logprobs", "advantages", "rollout_logprobs"] - for field in rl_fields: + # Handle loss side-channel fields. These need to be padded and sliced + # the same way as labels because they are token-aligned. + rl_field_dtypes = { + "target_tokens": torch.long, + "logprobs": torch.float, + "advantages": torch.float, + "rollout_logprobs": torch.float, + "teacher_ids": torch.long, + "teacher_cache_indices": torch.long, + "teacher_weights": torch.float, + "teacher_hidden_states": torch.float, + } + for field, dtype in rl_field_dtypes.items(): if field in batch: field_tensor = batch[field] if not isinstance(field_tensor, torch.Tensor): - # Determine dtype: logprobs/advantages are float, target_tokens is long - dtype = torch.float if field in ("logprobs", "advantages", "rollout_logprobs") else torch.long if isinstance(field_tensor, list): - if field_tensor and isinstance(field_tensor[0], list): - # Nested list - field_tensor = torch.tensor(field_tensor[0], dtype=dtype).unsqueeze(0) - else: - field_tensor = torch.tensor(field_tensor, dtype=dtype).unsqueeze(0) + field_tensor = torch.tensor(field_tensor, dtype=dtype) else: - field_tensor = torch.tensor(field_tensor, dtype=dtype).unsqueeze(0) + field_tensor = torch.tensor(field_tensor, dtype=dtype) + elif field == "teacher_hidden_states" and not torch.is_floating_point(field_tensor): + field_tensor = field_tensor.float() + + if field == "teacher_hidden_states": + if field_tensor.ndim == 2: + field_tensor = field_tensor.unsqueeze(0) + if field_tensor.ndim != 3: + raise ValueError( + f"teacher_hidden_states must have shape [batch, seq, hidden], got {field_tensor.shape}" + ) + seq_dim = 1 + elif field_tensor.ndim == 0: + field_tensor = field_tensor.unsqueeze(0) + seq_dim = -1 elif field_tensor.ndim == 1: field_tensor = field_tensor.unsqueeze(0) + seq_dim = -1 + else: + seq_dim = -1 # Determine pad value: IGNORE_INDEX for target_tokens, 0 for others pad_value = IGNORE_INDEX if field == "target_tokens" else 0.0 - field_tensor = self.sp_padding(field_tensor, dim=-1, pad_value=pad_value, pad_length=pad_length) + field_tensor = self.sp_padding(field_tensor, dim=seq_dim, pad_value=pad_value, pad_length=pad_length) if self.ringattn_size > 1: field_tensor = zigzag_reorder_packed_sequence( - field_tensor, original_position_ids, self.ringattn_size, dim=-1 + field_tensor, original_position_ids, self.ringattn_size, dim=seq_dim ) - batch[field] = self.sp_slice(field_tensor, dim=-1) + batch[field] = self.sp_slice(field_tensor, dim=seq_dim) # (Re)compute cu_seq_lens for flash attention. # For ring attention, position_ids has been zigzag-reordered: it has diff --git a/src/xorl/data/collators/tensor_collator.py b/src/xorl/data/collators/tensor_collator.py index caff02ca..921b094b 100644 --- a/src/xorl/data/collators/tensor_collator.py +++ b/src/xorl/data/collators/tensor_collator.py @@ -166,8 +166,11 @@ def _infer_dtype(self, key: str) -> torch.dtype: 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", "teacher_ids", "teacher_cache_indices"]: return torch.long + if key in ["teacher_weights"]: + return torch.float + # Let PyTorch infer for other fields return None # Will use default inference diff --git a/src/xorl/data/data_loader.py b/src/xorl/data/data_loader.py index ba2e3ff6..236d6b17 100644 --- a/src/xorl/data/data_loader.py +++ b/src/xorl/data/data_loader.py @@ -73,9 +73,11 @@ def __call__(self, features: Sequence[Dict[str, Any]]) -> List[Dict[str, Any]]: 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, but got {len(micro_batches)}" - ) + if len(micro_batches) != self.gradient_accumulation_steps: + raise RuntimeError( + f"Internal error: Expected {self.gradient_accumulation_steps} micro-batches, " + f"but got {len(micro_batches)}" + ) return micro_batches @@ -343,16 +345,19 @@ def build(self, verbose: bool = True) -> "DistributedDataloader": first_item = self.dataset[0] if isinstance(first_item, list): # New structure: list of dicts - assert len(first_item) > 0, "Dataset item is an empty list" - assert isinstance(first_item[0], dict), ( - f"Dataset items must be lists of dictionaries, but got list of {type(first_item[0]).__name__}. " - f"Each element in the list should be a dict with keys like 'input_ids', 'labels', etc." - ) + if len(first_item) == 0: + raise ValueError("Dataset item is an empty list") + if not isinstance(first_item[0], dict): + raise TypeError( + f"Dataset items must be lists of dictionaries, but got list of " + f"{type(first_item[0]).__name__}. Each element in the list should be a dict " + f"with keys like 'input_ids', 'labels', etc." + ) elif isinstance(first_item, dict): # Old structure: single dict (backward compatibility) pass else: - raise AssertionError( + raise TypeError( f"Dataset items must be either dict or list of dicts, but got {type(first_item).__name__}. " f"Each dataset item should be a dict (e.g., {{'input_ids': ..., 'labels': ...}}) " f"or a list of dicts (e.g., [{{'input_ids': ..., 'labels': ...}}, ...])." diff --git a/src/xorl/data/prepare/prepare_datasets.py b/src/xorl/data/prepare/prepare_datasets.py index 85c50c93..83c48b7b 100644 --- a/src/xorl/data/prepare/prepare_datasets.py +++ b/src/xorl/data/prepare/prepare_datasets.py @@ -100,10 +100,25 @@ def prepare_datasets( logger.info_rank0(f"Creating dummy dataset: {4096} samples x {seq_len} tokens") dataset = _create_dummy_dataset(seq_len=seq_len, seed=args.train.seed, vocab_size=len(tokenizer)) - if args.data.sample_packing_method and args.data.sample_packing_method != "none": - train_dataset = PackingDataset(args, tokenizer, dataset, split="train") - else: - train_dataset = dataset + def _build_train_dataset(): + if args.data.sample_packing_method and args.data.sample_packing_method != "none": + return PackingDataset(args, tokenizer, dataset, split="train") + return dataset + + # PackingDataset computes bins on rank 0 and caches them; non-zero + # ranks load from that cache. Gate with the same rank-0-first + + # barrier pattern used below so non-zero ranks don't race the cache. + is_distributed = dist.is_available() and dist.is_initialized() + is_rank_zero = not is_distributed or dist.get_rank() == 0 + + if is_rank_zero: + train_dataset = _build_train_dataset() + + if is_distributed: + dist.barrier() + + if is_distributed and not is_rank_zero: + train_dataset = _build_train_dataset() return train_dataset, None diff --git a/src/xorl/data/prepare/shared.py b/src/xorl/data/prepare/shared.py index 41e3f978..adb9d69c 100644 --- a/src/xorl/data/prepare/shared.py +++ b/src/xorl/data/prepare/shared.py @@ -512,8 +512,12 @@ def try_load_from_hub(args: Arguments, dataset_hash: str, split: str) -> Dataset token=args.data.hf_use_auth_token, ) return dataset[split] - except Exception: - LOG.info_rank0("Unable to find prepared dataset in HuggingFace Hub") + except Exception as e: + # `datasets.load_dataset` raises a wide variety of exception types + # (FileNotFoundError, requests.HTTPError, datasets.exceptions.*, etc.) + # so this remains a broad catch — but log the cause so failures are + # debuggable instead of silently turning into "not found". + LOG.info_rank0(f"Unable to find prepared dataset in HuggingFace Hub: {type(e).__name__}: {e}") return None diff --git a/src/xorl/distillation/__init__.py b/src/xorl/distillation/__init__.py new file mode 100644 index 00000000..d427e540 --- /dev/null +++ b/src/xorl/distillation/__init__.py @@ -0,0 +1,11 @@ +from xorl.distillation.teacher_cache import TeacherActivationCache, TeacherHeadManager, load_lm_head_weight +from xorl.distillation.teacher_store import TeacherHeadStore, prepare_lm_head_teacher_store + + +__all__ = [ + "TeacherHeadStore", + "TeacherActivationCache", + "TeacherHeadManager", + "load_lm_head_weight", + "prepare_lm_head_teacher_store", +] diff --git a/src/xorl/distillation/teacher_cache.py b/src/xorl/distillation/teacher_cache.py new file mode 100644 index 00000000..77e840ee --- /dev/null +++ b/src/xorl/distillation/teacher_cache.py @@ -0,0 +1,299 @@ +from __future__ import annotations + +import json +import os +from concurrent.futures import Future, ThreadPoolExecutor +from dataclasses import dataclass +from typing import Any, Dict, Mapping, Optional + +import torch +from safetensors import safe_open + +from xorl.distillation.teacher_store import ( + TeacherHeadShardView, + TeacherHeadStore, + is_teacher_store_entry, + load_lm_head_from_teacher_store, + teacher_store_manifest_path, +) + + +DEFAULT_LM_HEAD_KEY = "lm_head.weight" +DEFAULT_HIDDEN_KEY = "hidden_states" + + +def _load_safetensors_tensor(path: str, key: str) -> torch.Tensor: + with safe_open(path, framework="pt", device="cpu") as f: + if key not in f.keys(): + keys = list(f.keys()) + if len(keys) == 1: + return f.get_tensor(keys[0]) + raise KeyError(f"Could not find tensor key '{key}' in {path}. Available keys: {keys[:10]}") + return f.get_tensor(key) + + +def _load_from_model_dir(path: str, key: str) -> torch.Tensor: + index_path = os.path.join(path, "model.safetensors.index.json") + if os.path.exists(index_path): + with open(index_path) as f: + index = json.load(f) + shard_name = index.get("weight_map", {}).get(key) + if shard_name is None: + raise KeyError(f"Could not find tensor key '{key}' in {index_path}") + return _load_safetensors_tensor(os.path.join(path, shard_name), key) + + safetensors_path = os.path.join(path, "model.safetensors") + if os.path.exists(safetensors_path): + return _load_safetensors_tensor(safetensors_path, key) + + raise FileNotFoundError(f"No model.safetensors or model.safetensors.index.json found in {path}") + + +def _normalize_entry(entry: str | Mapping[str, Any], default_key: str) -> tuple[str, str]: + if isinstance(entry, str): + return entry, default_key + path = entry.get("path") or entry.get("model_path") or entry.get("weights_path") + if not path: + raise ValueError(f"Teacher entry must include a path: {entry}") + return path, entry.get("tensor_key", default_key) + + +def load_lm_head_weight( + entry: str | Mapping[str, Any], + tensor_key: str = DEFAULT_LM_HEAD_KEY, + teacher_id: int | str = 0, +) -> torch.Tensor: + """Load a teacher LM head from a local file or Hugging Face-style model directory.""" + if is_teacher_store_entry(entry): + return load_lm_head_from_teacher_store(entry, teacher_id=teacher_id).contiguous() + + path, key = _normalize_entry(entry, tensor_key) + if os.path.isdir(path): + tensor = _load_from_model_dir(path, key) + elif path.endswith(".safetensors"): + tensor = _load_safetensors_tensor(path, key) + else: + raise ValueError(f"Teacher LM head must be a safetensors file or model directory: {path}") + return tensor.contiguous() + + +def load_hidden_state_cache(entry: str | Mapping[str, Any], tensor_key: str = DEFAULT_HIDDEN_KEY) -> torch.Tensor: + """Load a teacher hidden-state cache tensor of shape [num_cached_tokens, hidden_dim].""" + path, key = _normalize_entry(entry, tensor_key) + if os.path.isdir(path): + safetensors_path = os.path.join(path, "hidden_states.safetensors") + if os.path.exists(safetensors_path): + tensor = _load_safetensors_tensor(safetensors_path, key) + else: + raise FileNotFoundError(f"No hidden_states.safetensors found in {path}") + elif path.endswith(".safetensors"): + tensor = _load_safetensors_tensor(path, key) + else: + raise ValueError(f"Teacher hidden cache must be a safetensors file or directory: {path}") + if tensor.ndim != 2: + raise ValueError(f"Teacher hidden cache must be rank 2 [tokens, hidden_dim], got shape {tuple(tensor.shape)}") + return tensor.contiguous() + + +@dataclass +class TeacherHeadManager: + """Holds teacher LM heads on CPU; promotes one at a time to device. + + Single-teacher device cache is intentional for the MVP: keeping all teacher + heads on device costs vocab_size * hidden_size per teacher (~600 MB for + Qwen3 vocab=151936, hidden=2048, fp32). Multi-teacher batches will thrash + on this re-upload; revisit with an LRU keyed on teacher_id when those + workloads land. + """ + + teacher_heads: Mapping[str, Any] + enable_async: bool = True + max_workers: int = 2 + + def __post_init__(self) -> None: + self.teacher_heads = {str(k): v for k, v in self.teacher_heads.items()} + self._cpu_cache: Dict[str, torch.Tensor] = {} + self._cpu_futures: Dict[str, Future] = {} + self._store_cache: Dict[str, TeacherHeadStore] = {} + self._executor: Optional[ThreadPoolExecutor] = ( + ThreadPoolExecutor(max_workers=max(1, int(self.max_workers)), thread_name_prefix="opd-teacher-head") + if self.enable_async + else None + ) + self._device_teacher_id: Optional[str] = None + self._device_tensor: Optional[torch.Tensor] = None + + @staticmethod + def _key(teacher_id: int | str | torch.Tensor) -> str: + return str(int(teacher_id)) if isinstance(teacher_id, torch.Tensor) else str(teacher_id) + + def _load_cpu(self, key: str) -> torch.Tensor: + cpu = load_lm_head_weight(self.teacher_heads[key], teacher_id=key) + if torch.cuda.is_available(): + cpu = cpu.pin_memory() + return cpu + + def has_sharded_head(self, teacher_id: int | str | torch.Tensor) -> bool: + key = self._key(teacher_id) + return key in self.teacher_heads and is_teacher_store_entry(self.teacher_heads[key]) + + def sharded_view( + self, + teacher_id: int | str | torch.Tensor, + device: torch.device | str, + dtype: Optional[torch.dtype] = None, + cache_cpu: bool = True, + cache_device: bool = False, + ) -> TeacherHeadShardView: + key = self._key(teacher_id) + if key not in self.teacher_heads: + raise KeyError(f"No teacher head configured for teacher_id={key}") + manifest_path = teacher_store_manifest_path(self.teacher_heads[key]) + if manifest_path is None: + raise ValueError(f"teacher_id={key} is not configured with a teacher store") + store_key = str(manifest_path) + if store_key not in self._store_cache: + self._store_cache[store_key] = TeacherHeadStore(manifest_path) + return TeacherHeadShardView( + store=self._store_cache[store_key], + teacher_id=key, + device=torch.device(device), + dtype=dtype, + cache_cpu=cache_cpu, + cache_device=cache_device, + ) + + def prefetch(self, teacher_id: int | str | torch.Tensor) -> None: + """Start loading a teacher head into pinned CPU memory.""" + key = self._key(teacher_id) + if key not in self.teacher_heads: + raise KeyError(f"No teacher head configured for teacher_id={key}") + if key in self._cpu_cache or key in self._cpu_futures: + return + if self._executor is None: + self._cpu_cache[key] = self._load_cpu(key) + else: + self._cpu_futures[key] = self._executor.submit(self._load_cpu, key) + + def _cpu_tensor(self, key: str) -> torch.Tensor: + if key in self._cpu_cache: + return self._cpu_cache[key] + if key in self._cpu_futures: + self._cpu_cache[key] = self._cpu_futures.pop(key).result() + return self._cpu_cache[key] + self._cpu_cache[key] = self._load_cpu(key) + return self._cpu_cache[key] + + def get( + self, teacher_id: int | str, device: torch.device | str, dtype: Optional[torch.dtype] = None + ) -> torch.Tensor: + key = self._key(teacher_id) + if key not in self.teacher_heads: + raise KeyError(f"No teacher head configured for teacher_id={key}") + cpu = self._cpu_tensor(key) + + target_device = torch.device(device) + if target_device.type == "cuda" and target_device.index is None and torch.cuda.is_available(): + target_device = torch.device("cuda", torch.cuda.current_device()) + needs_upload = ( + self._device_teacher_id != key + or self._device_tensor is None + or self._device_tensor.device != target_device + or (dtype is not None and self._device_tensor.dtype != dtype) + ) + if needs_upload: + self._device_tensor = None + self._device_teacher_id = key + self._device_tensor = cpu.to(device=target_device, dtype=dtype, non_blocking=True) + return self._device_tensor + + def close(self) -> None: + if self._executor is not None: + self._executor.shutdown(wait=False, cancel_futures=True) + self._executor = None + + +@dataclass +class TeacherActivationCache: + hidden_caches: Mapping[str, Any] | str + enable_async: bool = True + max_workers: int = 2 + + def __post_init__(self) -> None: + if not isinstance(self.hidden_caches, str): + self.hidden_caches = {str(k): v for k, v in self.hidden_caches.items()} + self._cpu_cache: Dict[str, torch.Tensor] = {} + self._cpu_futures: Dict[str, Future] = {} + self._executor: Optional[ThreadPoolExecutor] = ( + ThreadPoolExecutor(max_workers=max(1, int(self.max_workers)), thread_name_prefix="opd-teacher-hidden") + if self.enable_async + else None + ) + + def _entry_for_teacher(self, teacher_id: int | str) -> tuple[str, Any]: + key = str(int(teacher_id)) if isinstance(teacher_id, torch.Tensor) else str(teacher_id) + if isinstance(self.hidden_caches, str): + return key, self.hidden_caches + if key not in self.hidden_caches: + raise KeyError(f"No teacher hidden cache configured for teacher_id={key}") + return key, self.hidden_caches[key] + + def _load_cpu(self, key: str, entry: Any) -> torch.Tensor: + return load_hidden_state_cache(entry) + + def prefetch(self, teacher_id: int | str | torch.Tensor) -> None: + """Start loading a teacher hidden-state cache into CPU memory.""" + key, entry = self._entry_for_teacher(teacher_id) + if key in self._cpu_cache or key in self._cpu_futures: + return + if self._executor is None: + self._cpu_cache[key] = self._load_cpu(key, entry) + else: + self._cpu_futures[key] = self._executor.submit(self._load_cpu, key, entry) + + def _cpu_tensor(self, key: str, entry: Any) -> torch.Tensor: + if key in self._cpu_cache: + return self._cpu_cache[key] + if key in self._cpu_futures: + self._cpu_cache[key] = self._cpu_futures.pop(key).result() + return self._cpu_cache[key] + self._cpu_cache[key] = self._load_cpu(key, entry) + return self._cpu_cache[key] + + def get( + self, + teacher_id: int | str, + indices: torch.Tensor, + device: torch.device | str, + dtype: Optional[torch.dtype] = None, + ) -> torch.Tensor: + key, entry = self._entry_for_teacher(teacher_id) + cache = self._cpu_tensor(key, entry) + + flat_indices = indices.reshape(-1).to(device="cpu", dtype=torch.long) + if flat_indices.numel() > 0: + min_idx = flat_indices.min().item() + if min_idx < 0: + # Negative indices used to be silently clamped to 0, which masked + # producer bugs (off-by-one in teacher_cache_indices construction + # was found this way during the Countdown run). Fail loudly instead. + raise IndexError( + f"teacher_cache_indices contain negative value {min_idx} " + f"(teacher_id={key}); producer must emit non-negative indices" + ) + max_idx = flat_indices.max().item() + if max_idx >= cache.shape[0]: + raise IndexError( + f"teacher_cache_indices contain {max_idx}, " + f"but teacher_id={key} cache only has {cache.shape[0]} rows" + ) + gathered = cache.index_select(0, flat_indices) + gathered = gathered.view(*indices.shape, cache.shape[-1]) + # CPU cache is unpinned, so non_blocking=True would be a no-op. Drop the + # flag rather than misleading future readers. + return gathered.to(device=device, dtype=dtype) + + def close(self) -> None: + if self._executor is not None: + self._executor.shutdown(wait=False, cancel_futures=True) + self._executor = None diff --git a/src/xorl/distillation/teacher_store.py b/src/xorl/distillation/teacher_store.py new file mode 100644 index 00000000..11fd529e --- /dev/null +++ b/src/xorl/distillation/teacher_store.py @@ -0,0 +1,366 @@ +from __future__ import annotations + +import json +import os +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Dict, Iterator, Mapping, Optional + +import torch +from safetensors import safe_open +from safetensors.torch import save_file + + +TEACHER_STORE_MANIFEST = "manifest.json" +TEACHER_STORE_TYPE = "xorl_opd_teacher_store" +TEACHER_STORE_VERSION = 1 +DEFAULT_LM_HEAD_KEY = "lm_head.weight" + + +def _dtype_name(dtype: Any) -> str: + if isinstance(dtype, torch.dtype): + return str(dtype).removeprefix("torch.") + return str(dtype) + + +def _resolve(path: str | os.PathLike[str], base_dir: Path) -> Path: + p = Path(path) + return p if p.is_absolute() else base_dir / p + + +def _find_tensor_source(path: str | os.PathLike[str], key: str) -> tuple[Path, str]: + """Return the safetensors file and key containing ``key``. + + ``path`` can be a Hugging Face model directory, a single safetensors file, or + an already-materialized teacher-store manifest/directory. + """ + p = Path(path) + if p.is_dir(): + index_path = p / "model.safetensors.index.json" + if index_path.exists(): + index = json.loads(index_path.read_text(encoding="utf-8")) + shard_name = index.get("weight_map", {}).get(key) + if shard_name is None: + raise KeyError(f"Could not find tensor key '{key}' in {index_path}") + return p / shard_name, key + + safetensors_path = p / "model.safetensors" + if safetensors_path.exists(): + return safetensors_path, key + + if p.suffix == ".safetensors": + return p, key + + raise FileNotFoundError(f"No safetensors source found for key '{key}' at {p}") + + +@dataclass(frozen=True) +class TeacherHeadShard: + path: Path + tensor_key: str + start: int + end: int + + @property + def rows(self) -> int: + return self.end - self.start + + def load_cpu(self) -> torch.Tensor: + with safe_open(str(self.path), framework="pt", device="cpu") as f: + if self.tensor_key not in f.keys(): + keys = list(f.keys()) + if len(keys) != 1: + raise KeyError( + f"Could not find tensor key '{self.tensor_key}' in {self.path}. Available keys: {keys[:10]}" + ) + return f.get_tensor(keys[0]).contiguous() + return f.get_tensor(self.tensor_key).contiguous() + + +@dataclass(frozen=True) +class TeacherHeadSpec: + teacher_id: str + shape: tuple[int, int] + dtype: str + shards: tuple[TeacherHeadShard, ...] + + @property + def vocab_size(self) -> int: + return self.shape[0] + + @property + def hidden_size(self) -> int: + return self.shape[1] + + +class TeacherHeadStore: + """Row-sharded teacher LM-head store for OPD. + + The store keeps the teacher prediction head in vocab-row shards so OPD can + later stream only the shard needed by a vocab block instead of materializing + every teacher head on the GPU. + """ + + def __init__(self, manifest_path: str | os.PathLike[str]) -> None: + path = Path(manifest_path) + if path.is_dir(): + path = path / TEACHER_STORE_MANIFEST + self.manifest_path = path + self.root = path.parent + self.manifest = json.loads(path.read_text(encoding="utf-8")) + if self.manifest.get("type") != TEACHER_STORE_TYPE: + raise ValueError(f"{path} is not an XORL OPD teacher store manifest") + if int(self.manifest.get("version", 0)) != TEACHER_STORE_VERSION: + raise ValueError( + f"Unsupported teacher-store version {self.manifest.get('version')}; expected {TEACHER_STORE_VERSION}" + ) + self._cpu_cache: Dict[tuple[str, str, int, int, bool], torch.Tensor] = {} + + def teacher_ids(self) -> list[str]: + return list(self.manifest.get("teachers", {}).keys()) + + def head_spec(self, teacher_id: int | str) -> TeacherHeadSpec: + key = str(int(teacher_id)) if isinstance(teacher_id, torch.Tensor) else str(teacher_id) + teachers = self.manifest.get("teachers", {}) + if key not in teachers: + raise KeyError(f"No teacher_id={key} in teacher store {self.manifest_path}") + head = teachers[key]["lm_head"] + shards = tuple( + TeacherHeadShard( + path=_resolve(shard["path"], self.root), + tensor_key=shard.get("tensor_key", DEFAULT_LM_HEAD_KEY), + start=int(shard["start"]), + end=int(shard["end"]), + ) + for shard in head["shards"] + ) + return TeacherHeadSpec( + teacher_id=key, + shape=tuple(int(x) for x in head["shape"]), + dtype=head["dtype"], + shards=shards, + ) + + def iter_lm_head_shards(self, teacher_id: int | str) -> Iterator[tuple[int, int, torch.Tensor]]: + for shard in self.head_spec(teacher_id).shards: + yield shard.start, shard.end, shard.load_cpu() + + def load_shard_cpu( + self, + shard: TeacherHeadShard, + *, + pin_memory: bool = False, + cache: bool = True, + ) -> torch.Tensor: + key = (str(shard.path), shard.tensor_key, shard.start, shard.end, pin_memory) + if cache and key in self._cpu_cache: + return self._cpu_cache[key] + + tensor = shard.load_cpu() + if pin_memory: + tensor = tensor.pin_memory() + if cache: + self._cpu_cache[key] = tensor + return tensor + + def load_lm_head(self, teacher_id: int | str) -> torch.Tensor: + spec = self.head_spec(teacher_id) + pieces = [] + expected_start = 0 + for shard in spec.shards: + if shard.start != expected_start: + raise ValueError( + f"Teacher store {self.manifest_path} has non-contiguous lm_head shards: " + f"expected start {expected_start}, got {shard.start}" + ) + tensor = shard.load_cpu() + if tensor.shape[0] != shard.rows: + raise ValueError(f"Shard {shard.path} rows {tensor.shape[0]} do not match manifest rows {shard.rows}") + pieces.append(tensor) + expected_start = shard.end + out = torch.cat(pieces, dim=0).contiguous() + if tuple(out.shape) != spec.shape: + raise ValueError(f"Reconstructed lm_head shape {tuple(out.shape)} does not match manifest {spec.shape}") + return out + + +@dataclass(frozen=True) +class TeacherHeadShardView: + """Reusable view over a row-sharded teacher head. + + Shards are loaded lazily. CPU shard caching avoids repeated safetensors reads + across the multi-pass KL algorithm; optional device caching keeps one + teacher head resident for this view's forward/backward lifetime. + """ + + store: TeacherHeadStore + teacher_id: str + device: torch.device + dtype: Optional[torch.dtype] + cache_cpu: bool = True + cache_device: bool = False + _cpu_cache: Dict[tuple[str, str, int, int], torch.Tensor] = field( + default_factory=dict, + compare=False, + repr=False, + ) + _device_cache: Dict[tuple[str, str, int, int, str, str], torch.Tensor] = field( + default_factory=dict, + compare=False, + repr=False, + ) + + @property + def shape(self) -> tuple[int, int]: + return self.store.head_spec(self.teacher_id).shape + + def _cpu_tensor(self, shard: TeacherHeadShard) -> torch.Tensor: + key = (str(shard.path), shard.tensor_key, shard.start, shard.end) + if self.cache_cpu and key in self._cpu_cache: + return self._cpu_cache[key] + + pin_memory = self.device.type == "cuda" and torch.cuda.is_available() + cpu_tensor = self.store.load_shard_cpu( + shard, + pin_memory=pin_memory, + cache=self.cache_cpu, + ) + if self.cache_cpu: + self._cpu_cache[key] = cpu_tensor + return cpu_tensor + + def _device_tensor(self, shard: TeacherHeadShard, cpu_tensor: torch.Tensor) -> torch.Tensor: + dtype_name = str(self.dtype) if self.dtype is not None else str(cpu_tensor.dtype) + key = (str(shard.path), shard.tensor_key, shard.start, shard.end, str(self.device), dtype_name) + if self.cache_device and key in self._device_cache: + return self._device_cache[key] + + tensor = cpu_tensor.to(device=self.device, dtype=self.dtype, non_blocking=True) + if self.cache_device: + self._device_cache[key] = tensor + return tensor + + def iter_device_chunks(self, chunk_rows: int) -> Iterator[tuple[int, int, torch.Tensor]]: + if chunk_rows <= 0: + chunk_rows = self.shape[0] + for shard in self.store.head_spec(self.teacher_id).shards: + tensor = self._device_tensor(shard, self._cpu_tensor(shard)) + local_rows = shard.rows + for local_start in range(0, local_rows, chunk_rows): + local_end = min(local_start + chunk_rows, local_rows) + yield ( + shard.start + local_start, + shard.start + local_end, + tensor[local_start:local_end], + ) + + def clear_device_cache(self) -> None: + self._device_cache.clear() + + +def teacher_store_manifest_path(entry: str | os.PathLike[str] | Mapping[str, Any]) -> Optional[Path]: + explicit_store_path = False + if isinstance(entry, Mapping): + store_path = entry.get("store_path") or entry.get("teacher_store") or entry.get("manifest_path") + if store_path is None and entry.get("type") == TEACHER_STORE_TYPE: + store_path = entry.get("path") + if store_path is None: + return None + explicit_store_path = True + path = Path(store_path) + else: + path = Path(entry) + + if explicit_store_path: + return path / TEACHER_STORE_MANIFEST if path.suffix == "" else path + if path.is_dir() and (path / TEACHER_STORE_MANIFEST).exists(): + return path / TEACHER_STORE_MANIFEST + if path.is_file() and path.name == TEACHER_STORE_MANIFEST: + return path + return None + + +def is_teacher_store_entry(entry: str | os.PathLike[str] | Mapping[str, Any]) -> bool: + return teacher_store_manifest_path(entry) is not None + + +def load_lm_head_from_teacher_store( + entry: str | os.PathLike[str] | Mapping[str, Any], teacher_id: int | str +) -> torch.Tensor: + manifest_path = teacher_store_manifest_path(entry) + if manifest_path is None: + raise ValueError(f"Entry is not a teacher store: {entry}") + return TeacherHeadStore(manifest_path).load_lm_head(teacher_id) + + +def prepare_lm_head_teacher_store( + model_path: str | os.PathLike[str], + output_dir: str | os.PathLike[str], + *, + teacher_id: int | str = 0, + shard_rows: int = 32768, + tensor_key: str = DEFAULT_LM_HEAD_KEY, + force: bool = False, +) -> Path: + """Slice a teacher ``lm_head.weight`` into a row-sharded OPD teacher store.""" + if shard_rows <= 0: + raise ValueError(f"shard_rows must be positive, got {shard_rows}") + + output = Path(output_dir) + manifest_path = output / TEACHER_STORE_MANIFEST + if manifest_path.exists() and not force: + raise FileExistsError(f"{manifest_path} already exists; pass force=True to overwrite") + + source_path, source_key = _find_tensor_source(model_path, tensor_key) + output.mkdir(parents=True, exist_ok=True) + teacher_key = str(teacher_id) + teacher_dir = output / f"teacher_{teacher_key}" + teacher_dir.mkdir(parents=True, exist_ok=True) + + shards: list[Dict[str, Any]] = [] + with safe_open(str(source_path), framework="pt", device="cpu") as f: + if source_key not in f.keys(): + keys = list(f.keys()) + if len(keys) != 1: + raise KeyError( + f"Could not find tensor key '{source_key}' in {source_path}. Available keys: {keys[:10]}" + ) + source_key = keys[0] + tensor_slice = f.get_slice(source_key) + shape = tuple(int(x) for x in tensor_slice.get_shape()) + if len(shape) != 2: + raise ValueError(f"Teacher LM head must be rank 2 [vocab, hidden], got {shape}") + dtype = _dtype_name(tensor_slice.get_dtype()) + + for shard_idx, start in enumerate(range(0, shape[0], shard_rows)): + end = min(start + shard_rows, shape[0]) + tensor = tensor_slice[start:end].contiguous() + shard_name = f"lm_head_{shard_idx:05d}.safetensors" + shard_path = teacher_dir / shard_name + save_file({tensor_key: tensor}, str(shard_path)) + shards.append( + { + "path": str(shard_path.relative_to(output)), + "tensor_key": tensor_key, + "start": start, + "end": end, + } + ) + + manifest = { + "type": TEACHER_STORE_TYPE, + "version": TEACHER_STORE_VERSION, + "source": str(model_path), + "teachers": { + teacher_key: { + "lm_head": { + "tensor_key": tensor_key, + "shape": list(shape), + "dtype": dtype, + "shards": shards, + } + } + }, + } + manifest_path.write_text(json.dumps(manifest, indent=2, sort_keys=True) + "\n", encoding="utf-8") + return manifest_path diff --git a/src/xorl/distributed/fsdp2/__init__.py b/src/xorl/distributed/fsdp2/__init__.py index 4e9aa689..b5f649df 100644 --- a/src/xorl/distributed/fsdp2/__init__.py +++ b/src/xorl/distributed/fsdp2/__init__.py @@ -1 +1,2 @@ +from .bf16_a2a_reduce import BF16StochasticAllToAllReduceScatter from .clip_grad_norm import clip_grad_norm diff --git a/src/xorl/distributed/fsdp2/bf16_a2a_reduce.py b/src/xorl/distributed/fsdp2/bf16_a2a_reduce.py new file mode 100644 index 00000000..fb9d52c5 --- /dev/null +++ b/src/xorl/distributed/fsdp2/bf16_a2a_reduce.py @@ -0,0 +1,120 @@ +"""Custom FSDP2 reduce-scatter that exchanges in BF16 and sums in FP32. + +Implements the DeepSeek V4 §3.5.1 MoE gradient communication trick: + + 1. Stochastically round FP32 input → BF16 (unbiased). + 2. ``all_to_all_single`` the BF16 buffer across the reduce-scatter group. + 3. Sum the received per-rank chunks locally in FP32. + +Halves comm volume vs. native FP32 reduce-scatter while preserving FP32 +numerical robustness in the accumulator. NCCL ring/tree reduce-scatter +accumulates in the buffer's dtype during transit, so naively setting +FSDP ``reduce_dtype=bf16`` would accumulate partial sums in BF16's 8 +mantissa bits and accrue significant bias. Decoupling movement (BF16) from +accumulation (FP32) keeps the sum well-conditioned. + +Installed via ``FSDPModule.set_custom_reduce_scatter(...)`` on expert +FSDPModules; non-expert modules continue to use the default reduce-scatter. + +**Preconditions on the wrapping FSDPModule (enforced at install time in +``parallelize_model_fsdp2``; do NOT install this hook without them):** + + * ``mp_policy.reduce_dtype == torch.float32``. The hook accepts only FP32 + input — it relies on FSDP allocating the reduce-scatter buffer in FP32 + so that ``input_tensor`` arrives un-quantized; the BF16 transit is + internal to the hook. + * ``gradient_divide_factor == 1.0``. With factor=None FSDP enables a + ``predivide_factor`` that is applied to ``input_tensor`` *before* this + hook is invoked, and FSDP skips the postdivide because the reduce is + custom. The hook would then sum predivided values without compensation, + silently under-weighting gradients. The codebase calls + ``set_gradient_divide_factor(1.0)`` on every FSDPModule before installing + the hook (see ``torch_parallelize.py``); this hook will refuse to install + if that invariant is broken. +""" + +from typing import Optional, Sequence, Union + +import torch +import torch.distributed as dist +from torch.distributed.distributed_c10d import ReduceOp +from torch.distributed.fsdp._fully_shard._fsdp_api import ReduceScatter, _ReduceOp + +from xorl.optim.stochastic_round import stochastic_round_to_bf16 + + +def _canonical_reduce_op(op: _ReduceOp) -> ReduceOp.RedOpType: + """Return the underlying RedOpType for FSDP's wrapped or raw reduce op. + + Unknown ops are returned unchanged so the caller can reject them. + """ + op_type = getattr(op, "op", op) + op_name = getattr(op_type, "name", None) + if op_name == "SUM": + return ReduceOp.SUM + if op_name == "AVG": + return ReduceOp.AVG + if op_name == "PREMUL_SUM": + # FSDP emits _make_nccl_premul_sum(1 / gradient_divide_factor) when a + # gradient_divide_factor is set. PREMUL_SUM is only equivalent to SUM + # when that factor is 1.0; the install-time check in + # ``parallelize_model_fsdp2`` (see ``torch_parallelize.py``) that + # rejects gradient_divide_factor != 1.0 is load-bearing — without it, + # this branch silently drops the premultiply scalar and under-weights + # gradients. Do not relax that check without also inspecting the + # premul factor here. + return ReduceOp.SUM + return op_type + + +class BF16StochasticAllToAllReduceScatter(ReduceScatter): + """ReduceScatter: stochastic-round FP32→BF16, all-to-all, FP32 local sum.""" + + def allocate( + self, + size: Sequence[Union[int, torch.SymInt]], + *, + dtype: torch.dtype, + device: torch.device, + ) -> torch.Tensor: + return torch.empty(*size, dtype=dtype, device=device) + + def __call__( + self, + output_tensor: torch.Tensor, + input_tensor: torch.Tensor, + group: dist.ProcessGroup, + op: _ReduceOp, + async_op: bool = False, + ) -> Optional[dist.Work]: + if async_op: + raise NotImplementedError("BF16StochasticAllToAllReduceScatter does not support async_op=True") + op_type = _canonical_reduce_op(op) + if op_type != ReduceOp.SUM and op_type != ReduceOp.AVG: + raise NotImplementedError( + f"BF16StochasticAllToAllReduceScatter requires SUM or AVG op, got {op} ({op_type})" + ) + if input_tensor.dtype != torch.float32: + raise ValueError( + "BF16StochasticAllToAllReduceScatter requires FP32 input " + f"(set FSDP reduce_dtype=fp32), got {input_tensor.dtype}" + ) + + world_size = dist.get_world_size(group) + total_numel = input_tensor.numel() + if total_numel % world_size != 0: + raise ValueError( + f"Input numel {total_numel} not divisible by world_size {world_size}; " + "FSDP should already pad before calling this hook." + ) + chunk_numel = total_numel // world_size + + in_bf16 = stochastic_round_to_bf16(input_tensor) + out_bf16 = torch.empty_like(in_bf16) + dist.all_to_all_single(out_bf16, in_bf16, group=group) + + summed = out_bf16.view(world_size, chunk_numel).to(torch.float32).sum(dim=0) + if op_type == ReduceOp.AVG: + summed.div_(world_size) + output_tensor.copy_(summed) + return None diff --git a/src/xorl/distributed/fsdp2/clip_grad_norm.py b/src/xorl/distributed/fsdp2/clip_grad_norm.py index ef695722..e2d44803 100644 --- a/src/xorl/distributed/fsdp2/clip_grad_norm.py +++ b/src/xorl/distributed/fsdp2/clip_grad_norm.py @@ -32,9 +32,12 @@ def clip_grad_norm( reduce_groups = [] if fsdp_group is not None: + # get_world_size raises RuntimeError if the process group hasn't been + # initialized; treat that as a single-process group rather than a hard + # failure. Other exceptions (TypeError on bad argument, etc.) propagate. try: fsdp_world = dist.get_world_size(fsdp_group) - except Exception: + except RuntimeError: fsdp_world = 1 if fsdp_world > 1: reduce_groups.append(("fsdp", fsdp_group)) diff --git a/src/xorl/distributed/moe/deepep.py b/src/xorl/distributed/moe/deepep.py index 100d834e..b718e231 100644 --- a/src/xorl/distributed/moe/deepep.py +++ b/src/xorl/distributed/moe/deepep.py @@ -17,6 +17,11 @@ import torch import torch.distributed as dist +from ...utils import logging + + +logger = logging.get_logger(__name__) + try: import deep_ep @@ -49,6 +54,8 @@ def get_hidden_bytes(x: torch.Tensor) -> int: # --------------------------------------------------------------------------- _pending_combine_event: Optional["EventOverlap"] = None +_ALLOW_UNSAFE_ASYNC_COMBINE: bool = _os.environ.get("XORL_DEEPEP_UNSAFE_ASYNC_COMBINE", "0") == "1" + def sync_pending_combine(): """Wait for any pending async combine to complete. @@ -558,11 +565,23 @@ def tokens_post_combine( """Unpermute expert outputs and combine back to original ranks. Args: - async_combine: If True, combine runs asynchronously on the comm stream. - The output tensor data is NOT valid on the default stream until - ``sync_pending_combine()`` is called. The next call to - ``token_pre_dispatch()`` automatically syncs. + async_combine: Request asynchronous combine on the comm stream so the + wait can overlap with the next layer's compute. Honored only when + ``XORL_DEEPEP_UNSAFE_ASYNC_COMBINE=1`` is set; otherwise forced to + False because the combined tensor is consumed by the rest of the + transformer block on the default stream before the next dispatch, + and skipping ``event.current_stream_wait()`` here races that + consumer. """ + if async_combine and not _ALLOW_UNSAFE_ASYNC_COMBINE: + logger.warning_once( + "deepep_async_combine=True ignored: the combined tensor is consumed by the " + "transformer block on the default stream before the next DeepEP dispatch, so " + "deferring the sync races downstream compute. Set " + "XORL_DEEPEP_UNSAFE_ASYNC_COMBINE=1 to opt in anyway." + ) + async_combine = False + if _DEEPEP_PROFILE: return _tokens_post_combine_profiled(buffer, expert_output, ctx, async_combine) diff --git a/src/xorl/distributed/offloading.py b/src/xorl/distributed/offloading.py index dad2e69e..3716cbc0 100644 --- a/src/xorl/distributed/offloading.py +++ b/src/xorl/distributed/offloading.py @@ -1,6 +1,6 @@ import enum from contextlib import nullcontext -from typing import Tuple, Union +from typing import Optional, Tuple, Union import torch from torch.autograd.graph import saved_tensors_hooks @@ -56,9 +56,10 @@ def unpack_from_cpu(packed: Tuple[OffloadPolicy, torch.device, torch.Tensor]) -> def build_activation_offloading_context( enable_activation_offload: bool = False, enable_gradient_checkpointing: bool = False, - activation_gpu_limit: float = 0.0, + activation_gpu_limit: Optional[float] = 0.0, ) -> Tuple[Union["saved_tensors_hooks", "nullcontext"], Union["saved_tensors_hooks", "nullcontext"]]: model_fwd_context, model_bwd_context = nullcontext(), nullcontext() + activation_gpu_limit = 0.0 if activation_gpu_limit is None else activation_gpu_limit if enable_activation_offload: # pin_memory=False since CachingHostAllocator caches pinned memory aggressively. # torch._C._host_emptyCache() can be used after version 2.5. diff --git a/src/xorl/distributed/parallel_state.py b/src/xorl/distributed/parallel_state.py index 3f9c7dd4..fd37a433 100644 --- a/src/xorl/distributed/parallel_state.py +++ b/src/xorl/distributed/parallel_state.py @@ -207,14 +207,14 @@ def loss_group(self) -> Optional["ProcessGroup"]: All ranks that compute partial losses on different data/sequence shards. """ if self.device_mesh is not None: - return self.device_mesh.get_group("dp_sp") + return self.device_mesh.get_group(self._resolve_mesh_name("dp_sp")) return self.fsdp_group @property @requires_mesh def loss_mesh(self) -> "DeviceMesh": """Device mesh for loss reduction (dp_replicate x dp_shard x ulysses x ring).""" - return self.device_mesh["dp_sp"] + return self.device_mesh[self._resolve_mesh_name("dp_sp")] @property def loss_size(self) -> int: diff --git a/src/xorl/distributed/sequence_parallel/ring_attention.py b/src/xorl/distributed/sequence_parallel/ring_attention.py index adceb0ce..6a533903 100644 --- a/src/xorl/distributed/sequence_parallel/ring_attention.py +++ b/src/xorl/distributed/sequence_parallel/ring_attention.py @@ -84,11 +84,17 @@ def _all_gather_kv( ) -> Tuple[List[Tensor], List[Tensor]]: """All-gather KV from all ring ranks. - Stacks K and V into a single buffer, performs one all_gather, then splits. + Fast path (k.shape == v.shape): stacks K and V into a single buffer and + does one all_gather. + + MLA path (k.shape != v.shape, e.g. DeepSeek-V3 / Kimi with K having + ``qk_nope_head_dim + qk_rope_head_dim`` and V having ``v_head_dim``): + falls back to two separate all_gathers, mirroring + ``_gather_seq_scatter_kv`` in ``strategy.py``. Args: - k: local key tensor [B, S_local, H_kv, D] or [total_local, H_kv, D] - v: local value tensor, same shape as k + k: local key tensor [B, S_local, H_kv, D_k] or [total_local, H_kv, D_k] + v: local value tensor, same layout as k (D_v may differ from D_k for MLA) ringattn_group: ring attention process group Returns: @@ -96,10 +102,28 @@ def _all_gather_kv( """ ringattn_size = dist.get_world_size(ringattn_group) - # Pack K,V into single buffer: [2, *k.shape] - kv_local = torch.stack([k, v], dim=0).contiguous() + # MLA / mismatched head_dim fallback. + if k.shape != v.shape: + k_local = k.contiguous() + v_local = v.contiguous() + k_gathered = torch.empty( + (ringattn_size * k_local.shape[0], *k_local.shape[1:]), + dtype=k_local.dtype, + device=k_local.device, + ) + v_gathered = torch.empty( + (ringattn_size * v_local.shape[0], *v_local.shape[1:]), + dtype=v_local.dtype, + device=v_local.device, + ) + dist.all_gather_into_tensor(k_gathered, k_local, group=ringattn_group) + dist.all_gather_into_tensor(v_gathered, v_local, group=ringattn_group) + k_list = [chunk.contiguous() for chunk in k_gathered.chunk(ringattn_size, dim=0)] + v_list = [chunk.contiguous() for chunk in v_gathered.chunk(ringattn_size, dim=0)] + return k_list, v_list - # All-gather: [2 * ringattn_size, *k.shape] + # Fast path: pack K,V into single buffer [2, *k.shape] and all_gather once. + kv_local = torch.stack([k, v], dim=0).contiguous() kv_gathered = torch.empty( (ringattn_size * kv_local.shape[0], *kv_local.shape[1:]), dtype=kv_local.dtype, diff --git a/src/xorl/distributed/sequence_parallel/strategy.py b/src/xorl/distributed/sequence_parallel/strategy.py index 96342f50..188a8c04 100644 --- a/src/xorl/distributed/sequence_parallel/strategy.py +++ b/src/xorl/distributed/sequence_parallel/strategy.py @@ -19,7 +19,6 @@ import torch import torch.distributed as dist -from ...models.layers.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 @@ -59,6 +58,27 @@ def _scale_cu_seqlens_for_ringattn(kwargs, ringattn_group): return cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k +def _gather_seq_scatter_kv(k, v, *, seq_dim: int, head_dim: int, group): + """All-to-all K/V together when shapes match, otherwise dispatch separately.""" + if k.shape != v.shape: + k = gather_seq_scatter_heads(k, seq_dim=seq_dim, head_dim=head_dim, group=group) + v = gather_seq_scatter_heads(v, seq_dim=seq_dim, head_dim=head_dim, group=group) + return k.contiguous(), v.contiguous() + + if k.ndim == 3: + kv = torch.stack([k, v], dim=2) # [S, H_kv, 2, D] + kv = kv.reshape(k.size(0), 2 * k.size(1), k.size(2)) # [S, 2*H_kv, D] + kv = gather_seq_scatter_heads(kv, seq_dim=seq_dim, head_dim=head_dim, group=group) + kv = kv.reshape(kv.size(0), kv.size(1) // 2, 2, kv.size(2)) + return kv[:, :, 0, :].contiguous(), kv[:, :, 1, :].contiguous() + + 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 = gather_seq_scatter_heads(kv, seq_dim=seq_dim, head_dim=head_dim, group=group) + kv = kv.reshape(kv.size(0), kv.size(1), kv.size(2) // 2, 2, kv.size(3)) + return kv[:, :, :, 0, :].contiguous(), kv[:, :, :, 1, :].contiguous() + + # ------------------------------------------------------------------ # # Base class # ------------------------------------------------------------------ # @@ -177,26 +197,14 @@ def project_qkv(self, module, hidden_states, position_embeddings): # Q all-to-all (separate — Q has different head count from K/V in GQA) q = gather_seq_scatter_heads(q, seq_dim=0, head_dim=1, group=self.group) - # Fused KV all-to-all: interleave K/V heads into one tensor, single a2a - kv = torch.stack([k, v], dim=2) # [S, H_kv, 2, D] - kv = kv.reshape(k.size(0), 2 * k.size(1), k.size(2)) # [S, 2*H_kv, D] - kv = gather_seq_scatter_heads(kv, seq_dim=0, head_dim=1, group=self.group) - kv = kv.reshape(kv.size(0), kv.size(1) // 2, 2, kv.size(2)) # [S_full, H_kv/SP, 2, D] - k = kv[:, :, 0, :].contiguous() - v = kv[:, :, 1, :].contiguous() + k, v = _gather_seq_scatter_kv(k, v, seq_dim=0, head_dim=1, group=self.group) q = q.unsqueeze(0) k = k.unsqueeze(0) v = v.unsqueeze(0) else: 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 = 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] - k = kv[:, :, :, 0, :].contiguous() - v = kv[:, :, :, 1, :].contiguous() + k, v = _gather_seq_scatter_kv(k, v, seq_dim=1, head_dim=2, group=self.group) return q, k, v @@ -247,21 +255,13 @@ def project_qkv(self, module, hidden_states, position_embeddings): k = k.squeeze(0) v = v.squeeze(0) q = gather_seq_scatter_heads(q, seq_dim=0, head_dim=1, group=self.group) - kv = torch.stack([k, v], dim=2).reshape(k.size(0), 2 * k.size(1), k.size(2)) - kv = gather_seq_scatter_heads(kv, seq_dim=0, head_dim=1, group=self.group) - kv = kv.reshape(kv.size(0), kv.size(1) // 2, 2, kv.size(2)) - k = kv[:, :, 0, :].contiguous() - v = kv[:, :, 1, :].contiguous() + k, v = _gather_seq_scatter_kv(k, v, seq_dim=0, head_dim=1, group=self.group) q = q.unsqueeze(0) k = k.unsqueeze(0) v = v.unsqueeze(0) else: q = gather_seq_scatter_heads(q, seq_dim=1, head_dim=2, group=self.group) - kv = torch.stack([k, v], dim=3).reshape(k.size(0), k.size(1), 2 * k.size(2), k.size(3)) - 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)) - k = kv[:, :, :, 0, :].contiguous() - v = kv[:, :, :, 1, :].contiguous() + k, v = _gather_seq_scatter_kv(k, v, seq_dim=1, head_dim=2, group=self.group) return q, k, v # Squeeze to 2D [S, H] — async autograd Function requires 2D (manual matmul in backward) @@ -309,6 +309,9 @@ def project_qkv(self, module, hidden_states, position_embeddings): k = k.unsqueeze(0) v = v.unsqueeze(0) + # Lazy-imported to break an import cycle with xorl.models.layers.attention + from ...models.layers.rope import apply_rotary_pos_emb # noqa: PLC0415 + cos, sin = position_embeddings # full-length S_full (NOT sliced) q, k = apply_rotary_pos_emb(q, k, cos, sin) return q, k, v diff --git a/src/xorl/distributed/sync_padding.py b/src/xorl/distributed/sync_padding.py index 4953df13..f7111e5b 100644 --- a/src/xorl/distributed/sync_padding.py +++ b/src/xorl/distributed/sync_padding.py @@ -100,11 +100,13 @@ def synchronize_micro_batch_padding( n_keys = len(all_keys) for i, mb in enumerate(micro_batches): offset = i * n_keys + max_sequence_target = 0 # 2D keys for j, key in enumerate(keys_2d): target = target_lens[offset + j] current = mb[key].shape[-1] + max_sequence_target = max(max_sequence_target, target) if current >= target: continue pad_amt = target - current @@ -119,3 +121,22 @@ def synchronize_micro_batch_padding( pad_amt = target - current # F.pad pads from last dim backwards: (H_left, H_right, S_left, S_right) mb[key] = F.pad(mb[key], (0, 0, 0, pad_amt), value=_3D_PAD_VALUES[key]) + + # Keep FA varlen metadata consistent with the padded tensor length. + # Shorter DP ranks are padded to the global max sequence length above; + # FA3 expects q/k/v.shape[0] to match cu_seq_lens[-1]. Grow the final + # sequence segment to include sync padding rather than adding a new + # padding segment. + for key in ("cu_seq_lens_q", "cu_seq_lens_k"): + if key in mb and isinstance(mb[key], torch.Tensor) and max_sequence_target > 0: + if mb[key][-1] < max_sequence_target: + mb[key] = mb[key].clone() + mb[key][-1] = max_sequence_target + + for max_key, cu_key in (("max_length_q", "cu_seq_lens_q"), ("max_length_k", "cu_seq_lens_k")): + if cu_key in mb and isinstance(mb[cu_key], torch.Tensor): + new_max = int(mb[cu_key].diff().max().item()) + if max_key in mb: + mb[max_key] = max(int(mb[max_key]), new_max) + else: + mb[max_key] = new_max diff --git a/src/xorl/distributed/torch_parallelize.py b/src/xorl/distributed/torch_parallelize.py index e92dbca1..92b6af70 100644 --- a/src/xorl/distributed/torch_parallelize.py +++ b/src/xorl/distributed/torch_parallelize.py @@ -12,14 +12,15 @@ from torch.utils.checkpoint import noop_context_fn from xorl.distributed.checkpoint import CheckpointFunction -from xorl.distributed.fsdp2 import clip_grad_norm +from xorl.distributed.fsdp2 import BF16StochasticAllToAllReduceScatter, clip_grad_norm from xorl.distributed.parallel_state import get_parallel_state from xorl.distributed.pipeline_parallel import ( generate_llm_fqn_per_model_part, pipeline_module_split, ) from xorl.lora import LoraLinear -from xorl.models import all_ranks_load_weights, rank0_load_and_broadcast_weights +from xorl.models import all_ranks_load_weights, grouped_load_weights +from xorl.models.transformers.deepseek_v3.support import validate_deepseek_v3_tensor_parallelism from xorl.utils import logging from xorl.utils.device import get_device_type from xorl.utils.import_utils import is_torch_version_greater_than @@ -29,6 +30,7 @@ from torch.distributed._composable.fsdp import MixedPrecisionPolicy, fully_shard from torch.distributed.tensor.parallel import ( ColwiseParallel, + ParallelStyle, RowwiseParallel, parallelize_module, ) @@ -37,6 +39,36 @@ logger = logging.get_logger(__name__) +def _load_model_weights( + model: "nn.Module", + weights_path: str, + load_weights_mode: str, + weight_device: str, + dtensor_factory=None, +) -> None: + """Dispatch HF weight loading by mode. + + 'skip' is a no-op; weights are expected to come from a DCP checkpoint via + load_checkpoint_path. 'grouped' uses one reader per node for dense weights + and one per EP-FSDP group for experts, falling back to rank-0 loading when + those fanout groups are unavailable. 'all_ranks' has every rank read + independently. + """ + if load_weights_mode == "skip": + logger.info_rank0("Skipping HF weight loading (weights will be loaded from DCP checkpoint).") + return + if load_weights_mode == "grouped": + logger.info_rank0("Loading model weights with one reader per node (dense) + per EP-FSDP group (experts)...") + grouped_load_weights(model, weights_path, weight_device, dtensor_factory=dtensor_factory) + elif load_weights_mode == "all_ranks": + logger.info_rank0("Every rank reading weights from disk independently...") + all_ranks_load_weights(model, weights_path, weight_device, dtensor_factory=dtensor_factory) + else: + raise ValueError( + f"Unsupported load_weights_mode={load_weights_mode!r}. Expected one of: grouped, all_ranks, skip." + ) + + _TP_STYLE_MAP = { "colwise": ColwiseParallel if is_torch_version_greater_than("2.4") else None, "rowwise": RowwiseParallel if is_torch_version_greater_than("2.4") else None, @@ -71,19 +103,27 @@ def _build_tp_plan(model: "nn.Module") -> Dict[str, Any]: return plan -def _resolve_tp_style(style_str: str): - """Convert a string TP style to a PyTorch ParallelStyle object.""" - if style_str == "colwise_rep": +def _resolve_tp_style(style): + """Convert a string TP style to a PyTorch ``ParallelStyle`` instance. + + ``style`` may also already be a ``ParallelStyle`` instance — useful for + plans that need configuration the string shortcuts don't cover (e.g. + ``ColwiseParallel(input_layouts=...)``, custom user-defined styles). + Returned as-is in that case. + """ + if isinstance(style, ParallelStyle): + return style + if style == "colwise_rep": return ColwiseParallel(output_layouts=Replicate()) - elif style_str == "embedding": + elif style == "embedding": # Embedding: shard weight on vocab dim, replicated input/output # Weight [vocab, hidden] → Shard(0) [vocab/tp, hidden] per rank # Lookup → partial results → all-reduce → replicated output return RowwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()) - elif style_str in _TP_STYLE_MAP: - return _TP_STYLE_MAP[style_str]() + elif style in _TP_STYLE_MAP: + return _TP_STYLE_MAP[style]() else: - raise ValueError(f"Unknown TP style: {style_str}") + raise ValueError(f"Unknown TP style: {style}") def parallelize_model_fsdp2( @@ -93,6 +133,7 @@ def parallelize_model_fsdp2( basic_modules: Optional[List[str]] = None, pp_enabled: bool = False, reshard_after_forward: Optional[bool] = None, + moe_grad_reduce_mode: str = "reduce_scatter", **kwargs, ) -> "nn.Module": """ @@ -215,6 +256,7 @@ def _experts_shard_placement_fn(param): # | -- experts layer (apply fully_shard separately in order to shard across EP groups on the same EP rank instead of sharding globally) # | -- layers (declared in model.modules_to_ignore_in_mixed_precision) that need to apply fully_shard separately due to different mp policy as the decoder layer # (e.g., some models requires MoE TopK gate layer to have parameters in higher FP32 precision in forward). + fsdp_wrapped_experts: List["nn.Module"] = [] for layer_fqn, layer_mod, experts_mod in layer_pairs: # register all the FSDPModule inside this decoder layer for the convenience of manual prefetching configuration layer_mod._fsdp_modules = [] @@ -233,6 +275,7 @@ def _experts_shard_placement_fn(param): # shard expert fully_shard(experts_mod, **expert_fsdp_kwargs) layer_mod._fsdp_modules.append(experts_mod) + fsdp_wrapped_experts.append(experts_mod) # shard module that needs to ignore mixed precision control if mp_ignored_classes: for sub_mod in layer_mod.modules(): @@ -330,13 +373,67 @@ def _experts_shard_placement_fn(param): if isinstance(module, FSDPModule): module.set_gradient_divide_factor(1.0) + # Install custom reduce-scatter for MoE expert FSDP units. + # Native FSDP reduce-scatter performs accumulation in the buffer dtype during transit, + # so naively setting reduce_dtype=bf16 would corrupt the partial sum. The + # bf16_a2a_fp32_sum mode keeps the FSDP buffer at FP32 (set via mp_policy.reduce_dtype) + # but stochastically rounds to BF16 for the all-to-all and sums received chunks + # locally in FP32. Halves comm volume per DeepSeek V4 §3.5.1. + # + # Preconditions enforced below (see BF16StochasticAllToAllReduceScatter docstring): + # - mp_policy.reduce_dtype must be torch.float32 (the hook receives the FSDP + # reduce buffer pre-allocated in this dtype; rejecting at __call__ time would + # surface the error during the first backward, far from configuration). + # - gradient_divide_factor must be 1.0. With factor=None FSDP enables a + # predivide_factor that gets applied to the input *before* our hook sees it, + # and FSDP would skip the postdivide because we own the reduce; the result + # is silently under-weighted gradients. + if moe_grad_reduce_mode == "bf16_a2a_fp32_sum": + if fsdp_wrapped_experts: + for experts_mod in fsdp_wrapped_experts: + state = experts_mod._get_fsdp_state() + pg = state._fsdp_param_group + if pg is None: + continue + if pg.mp_policy.reduce_dtype != torch.float32: + raise ValueError( + "moe_grad_reduce_mode='bf16_a2a_fp32_sum' requires FSDP " + f"mp_policy.reduce_dtype=torch.float32, got {pg.mp_policy.reduce_dtype}. " + "Either keep enable_mixed_precision=True (which sets reduce_dtype=fp32) " + "or pass an explicit MixedPrecisionPolicy with reduce_dtype=torch.float32." + ) + if pg.gradient_divide_factor != 1.0: + raise ValueError( + "moe_grad_reduce_mode='bf16_a2a_fp32_sum' requires " + "gradient_divide_factor=1.0 on every expert FSDP unit " + f"(got {pg.gradient_divide_factor}). FSDP applies predivide before " + "the reduce-scatter hook and skips postdivide when the hook is custom, " + "which would silently under-weight expert gradients." + ) + for experts_mod in fsdp_wrapped_experts: + experts_mod.set_custom_reduce_scatter(BF16StochasticAllToAllReduceScatter()) + logger.info_rank0( + f"Installed BF16 stochastic-rounded all-to-all + FP32 local sum reduce-scatter " + f"on {len(fsdp_wrapped_experts)} expert FSDP units." + ) + else: + logger.warning_rank0( + "moe_grad_reduce_mode='bf16_a2a_fp32_sum' was set but no expert FSDP units " + "were wrapped (EP not enabled or experts use _skip_fsdp). Mode has no effect." + ) + elif moe_grad_reduce_mode != "reduce_scatter": + raise ValueError( + f"Unsupported moe_grad_reduce_mode: {moe_grad_reduce_mode!r}. " + "Expected 'reduce_scatter' or 'bf16_a2a_fp32_sum'." + ) + # Handle meta initialization for FSDP2 (fallback if pre-load not done) assert kwargs.get("init_device") == "meta", "Please use init_device: meta for FSDP2" weight_device = get_device_type() # skip_weight_loading: Used when caller will handle weight loading separately - # (e.g., FSDP2+LoRA where we broadcast from rank 0 after this function returns) + # after this function returns. if kwargs.get("skip_weight_loading"): logger.info_rank0("Skipping weight loading in parallelize_model_fsdp2 (caller will handle)") elif weights_path is None: @@ -350,13 +447,14 @@ def _experts_shard_placement_fn(param): m.reset_lora_parameters() else: logger.info_rank0("starting to load model weights...") - load_weights_mode = kwargs.get("load_weights_mode", "broadcast") - if load_weights_mode == "broadcast": - logger.info_rank0("Loading model weights from disk on rank0 then broadcasting to other ranks...") - rank0_load_and_broadcast_weights(model, weights_path, weight_device, dtensor_factory=distribute_tensor) - else: - logger.info_rank0("Every rank reading weights from disk independently...") - all_ranks_load_weights(model, weights_path, weight_device, dtensor_factory=distribute_tensor) + load_weights_mode = kwargs.get("load_weights_mode", "grouped") + _load_model_weights( + model, + weights_path, + load_weights_mode=load_weights_mode, + weight_device=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 @@ -438,6 +536,7 @@ def build_parallelize_model( basic_modules: Optional[List[str]] = None, pp_schedule: Optional[str] = None, reshard_after_forward: Optional[bool] = None, + moe_grad_reduce_mode: str = "reduce_scatter", **kwargs, ) -> "nn.Module": """ @@ -523,6 +622,7 @@ def build_parallelize_model( # TP (if enabled) if ps.tp_enabled: + validate_deepseek_v3_tensor_parallelism(model_part.config) # TP + LoRA is not currently supported if i == 0: if any(isinstance(m, LoraLinear) for m in model_part.modules()): @@ -549,15 +649,14 @@ def build_parallelize_model( if kwargs.get("init_device") == "meta" and weights_path is not None: 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 - ) - else: - all_ranks_load_weights( - model_part, weights_path, get_device_type(), dtensor_factory=distribute_tensor - ) + load_weights_mode = kwargs.get("load_weights_mode", "grouped") + _load_model_weights( + model_part, + weights_path, + load_weights_mode=load_weights_mode, + weight_device=get_device_type(), + dtensor_factory=distribute_tensor, + ) kwargs["skip_weight_loading"] = True # torch.compile (if enabled) @@ -593,6 +692,7 @@ def build_parallelize_model( basic_modules=basic_modules, pp_enabled=True, reshard_after_forward=reshard_after_forward, + moe_grad_reduce_mode=moe_grad_reduce_mode, **kwargs, ) elif ps.dp_mode == "ddp": @@ -668,6 +768,7 @@ def _reentrant_ckpt_with_kwargs(fn, *args, **kw): module._gradient_checkpointing_func = _make_wrapper(_orig) if parallel_state.tp_enabled: + validate_deepseek_v3_tensor_parallelism(model.config) # TP + LoRA is not currently supported if any(isinstance(m, LoraLinear) for m in model.modules()): raise NotImplementedError("Tensor parallelism + LoRA is not currently supported.") @@ -692,13 +793,14 @@ def _reentrant_ckpt_with_kwargs(fn, *args, **kw): # Load weights now so FSDP wraps materialized TP DTensors. if kwargs.get("init_device") == "meta" and weights_path is not None: 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 - ) - else: - all_ranks_load_weights(model, weights_path, get_device_type(), dtensor_factory=distribute_tensor) + load_weights_mode = kwargs.get("load_weights_mode", "grouped") + _load_model_weights( + model, + weights_path, + load_weights_mode=load_weights_mode, + weight_device=get_device_type(), + dtensor_factory=distribute_tensor, + ) # Mark weights as already loaded so FSDP path skips loading kwargs["skip_weight_loading"] = True @@ -731,6 +833,7 @@ def _reentrant_ckpt_with_kwargs(fn, *args, **kw): enable_mixed_precision=enable_mixed_precision, basic_modules=basic_modules, reshard_after_forward=reshard_after_forward, + moe_grad_reduce_mode=moe_grad_reduce_mode, **kwargs, ) elif parallel_state.dp_mode == "ddp": diff --git a/src/xorl/lora/__init__.py b/src/xorl/lora/__init__.py index 1a698376..752bad0a 100644 --- a/src/xorl/lora/__init__.py +++ b/src/xorl/lora/__init__.py @@ -6,13 +6,7 @@ """ from xorl.lora.modules import LoraLinear, LoraModule -from xorl.ops.group_gemm.kernel import ( - compute_lora_scaling, - get_lora_delta_weight_stacked, - init_lora_weights_stacked, - merge_lora_weights_stacked, - unmerge_lora_weights_stacked, -) +from xorl.ops.group_gemm.kernel import compute_lora_scaling _MAPPING_ATTRS = { @@ -70,11 +64,7 @@ "copy_weights_to_lora_experts", "mark_only_lora_as_trainable", "lora_state_dict", - "init_lora_weights_stacked", "compute_lora_scaling", - "merge_lora_weights_stacked", - "unmerge_lora_weights_stacked", - "get_lora_delta_weight_stacked", ] diff --git a/src/xorl/lora/modules/linear.py b/src/xorl/lora/modules/linear.py index c6c4772f..78917f14 100644 --- a/src/xorl/lora/modules/linear.py +++ b/src/xorl/lora/modules/linear.py @@ -68,6 +68,8 @@ def __init__( # LoRA-specific attributes self.r = r self.lora_alpha = lora_alpha + self.active_r = r + self.active_lora_alpha = lora_alpha self.scaling = lora_alpha / r # LoRA weights (trainable, float32 for numerical stability) @@ -151,6 +153,16 @@ def reset_lora_parameters(self) -> None: # LoRA B: zeros (ensures output starts unchanged) nn.init.zeros_(self.lora_B) + def set_runtime_lora_config(self, lora_rank: int, lora_alpha: int) -> None: + """Update the active LoRA slice used during forward/merge/export.""" + if lora_rank <= 0 or lora_rank > self.r: + raise ValueError(f"Active LoRA rank must be in [1, {self.r}], got {lora_rank}") + self.active_r = lora_rank + self.active_lora_alpha = lora_alpha + + def _active_scaling(self) -> float: + return self.active_lora_alpha / self.active_r + def forward(self, x: torch.Tensor) -> torch.Tensor: """ Forward pass with LoRA adaptation. @@ -167,11 +179,13 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: # LoRA path: A -> B -> scale # Compute in float32 for numerical stability x_lora = x.to(self.lora_A.dtype) + lora_A = self.lora_A[: self.active_r] + lora_B = self.lora_B[:, : self.active_r] # x_lora @ lora_A.T @ lora_B.T * scaling # = F.linear(F.linear(x_lora, lora_A), lora_B) * scaling - lora_out = F.linear(F.linear(x_lora, self.lora_A), self.lora_B) - lora_out = lora_out * self.scaling + lora_out = F.linear(F.linear(x_lora, lora_A), lora_B) + lora_out = lora_out * self._active_scaling() # Add LoRA output to result (cast back to result dtype) return result + lora_out.to(result.dtype) @@ -183,7 +197,9 @@ def get_delta_weight(self) -> torch.Tensor: Returns: Delta weight tensor: lora_B @ lora_A * scaling """ - return (self.lora_B @ self.lora_A) * self.scaling + lora_A = self.lora_A[: self.active_r] + lora_B = self.lora_B[:, : self.active_r] + return (lora_B @ lora_A) * self._active_scaling() def merge_weights(self) -> None: """ @@ -198,12 +214,14 @@ def merge_weights(self) -> None: """ with torch.no_grad(): delta_weight = self.get_delta_weight() - self.weight.add_(delta_weight.to(self.weight.dtype)) + merged = self.weight.to(torch.float32) + delta_weight + self.weight.data.copy_(merged.to(self.weight.dtype)) self.reset_lora_parameters() def extra_repr(self) -> str: # Extend nn.Linear's extra_repr with LoRA-specific info return ( f"in_features={self.in_features}, out_features={self.out_features}, " - f"bias={self.bias is not None}, r={self.r}, lora_alpha={self.lora_alpha}" + f"bias={self.bias is not None}, r={self.active_r}, max_r={self.r}, " + f"lora_alpha={self.active_lora_alpha}" ) diff --git a/src/xorl/lora/utils.py b/src/xorl/lora/utils.py index bf601e08..6d7bf36e 100644 --- a/src/xorl/lora/utils.py +++ b/src/xorl/lora/utils.py @@ -8,7 +8,8 @@ import logging import os import re -from typing import Dict, Iterator, List, Optional, Tuple +from dataclasses import dataclass +from typing import Dict, Iterable, Iterator, List, Optional, Tuple import torch import torch.distributed as dist @@ -30,6 +31,36 @@ "qwen": ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"], "mistral": ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"], "gemma": ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"], + "deepseek_v3": [ + "q_a_proj", + "q_b_proj", + "kv_a_proj_with_mqa", + "kv_b_proj", + "o_proj", + "gate_proj", + "up_proj", + "down_proj", + ], + "kimi_k2": [ + "q_a_proj", + "q_b_proj", + "kv_a_proj_with_mqa", + "kv_b_proj", + "o_proj", + "gate_proj", + "up_proj", + "down_proj", + ], + "kimi_k25": [ + "q_a_proj", + "q_b_proj", + "kv_a_proj_with_mqa", + "kv_b_proj", + "o_proj", + "gate_proj", + "up_proj", + "down_proj", + ], } _PEFT_BASE_MODEL_PREFIX = "base_model.model." @@ -37,6 +68,34 @@ _MOE_PEFT_LORA_PATTERN = re.compile( r"(.*)\.mlp\.experts\.(shared|\d+)\.(gate_proj|up_proj|down_proj)\.lora_(A|B)\.weight$" ) +_MOE_SGLANG_SHARED_OUTER_PATTERN = re.compile(r"(.*)\.mlp\.experts\.(w1|w2|w3)\.lora_(A|B)\.weight$") + +# SGLang shared_outer format uses w1/w2/w3 slots for gate/down/up projections. +_PROJ_TO_SGLANG_W = {"gate_proj": "w1", "down_proj": "w2", "up_proj": "w3"} +_SGLANG_W_TO_PROJ = {v: k for k, v in _PROJ_TO_SGLANG_W.items()} + +LORA_EXPORT_FORMATS = ("peft", "sglang_shared_outer") + + +@dataclass(frozen=True) +class LoraTensorShardSpec: + """Shard metadata needed to map a full LoRA tensor onto the current rank.""" + + dim: int + index: int + size: int + + +def _get_default_target_modules(model: nn.Module) -> List[str]: + config = getattr(model, "config", None) + model_type = getattr(config, "model_type", None) + if model_type in DEFAULT_TARGET_MODULES: + return list(DEFAULT_TARGET_MODULES[model_type]) + if model_type is not None: + for family, targets in DEFAULT_TARGET_MODULES.items(): + if family in model_type: + return list(targets) + return ["q_proj", "k_proj", "v_proj", "o_proj"] def _get_submodule(model: nn.Module, target: str) -> Tuple[nn.Module, str]: @@ -160,7 +219,7 @@ def inject_lora_into_model( ... ) """ if target_modules is None: - target_modules = ["q_proj", "k_proj", "v_proj", "o_proj"] + target_modules = _get_default_target_modules(model) # Find all matching modules target_paths = _find_target_modules(model, target_modules) @@ -279,6 +338,75 @@ def _gather_ep_tensor(tensor: torch.Tensor, spec_info) -> torch.Tensor: return torch.cat(gathered, dim=shard_dim) +def _lora_rank_dim(param_name: str) -> Optional[int]: + """Return the rank dimension for xorl LoRA parameter names.""" + if param_name == "lora_A": + return 0 + if param_name == "lora_B": + return 1 + if param_name.endswith("_lora_A"): + return -1 + if param_name.endswith("_lora_B"): + return 1 + return None + + +def _active_lora_rank_slices(model: nn.Module) -> Dict[str, Tuple[int, int]]: + """Map LoRA parameter FQNs to (rank_dim, active_rank) from live modules.""" + rank_slices: Dict[str, Tuple[int, int]] = {} + for module_name, module in model.named_modules(): + active_rank = getattr(module, "active_r", None) + if active_rank is None: + continue + active_rank = int(active_rank) + if active_rank <= 0: + raise ValueError(f"Active LoRA rank must be positive, got {active_rank}") + + prefix = f"{module_name}." if module_name else "" + for local_name, _ in module.named_parameters(recurse=False): + rank_dim = _lora_rank_dim(local_name) + if rank_dim is not None: + rank_slices[f"{prefix}{local_name}"] = (rank_dim, active_rank) + return rank_slices + + +def _slice_lora_tensor_to_rank(tensor: torch.Tensor, rank_dim: int, active_rank: int) -> torch.Tensor: + dim = rank_dim if rank_dim >= 0 else tensor.dim() + rank_dim + if dim < 0 or dim >= tensor.dim() or tensor.shape[dim] <= active_rank: + return tensor + return tensor.narrow(dim, 0, active_rank).contiguous() + + +def slice_lora_state_dict_to_active_rank( + model: nn.Module, + lora_state_dict: Dict[str, torch.Tensor], +) -> Dict[str, torch.Tensor]: + """Slice LoRA state tensors to each module's active runtime rank.""" + rank_slices = _active_lora_rank_slices(model) + if not rank_slices: + return lora_state_dict + + sliced_state_dict: Dict[str, torch.Tensor] = {} + for name, tensor in lora_state_dict.items(): + rank_slice = rank_slices.get(name) + if rank_slice is None: + sliced_state_dict[name] = tensor + continue + rank_dim, active_rank = rank_slice + sliced_state_dict[name] = _slice_lora_tensor_to_rank(tensor, rank_dim, active_rank) + return sliced_state_dict + + +def _first_active_lora_config(model: nn.Module) -> Tuple[Optional[int], Optional[int]]: + """Return the first live module's active rank/alpha, if present.""" + for module in model.modules(): + active_rank = getattr(module, "active_r", None) + active_alpha = getattr(module, "active_lora_alpha", None) + if active_rank is not None: + return int(active_rank), None if active_alpha is None else int(active_alpha) + return None, None + + def get_lora_state_dict( model: nn.Module, prefix: str = "", @@ -300,6 +428,7 @@ def get_lora_state_dict( """ fqn2spec_info = getattr(model, "_fqn2spec_info", None) + rank_slices = _active_lora_rank_slices(model) lora_state_dict = {} for name, param in model.named_parameters(): @@ -319,6 +448,11 @@ def get_lora_state_dict( if spec_info is not None: tensor = _gather_ep_tensor(tensor, spec_info) + rank_slice = rank_slices.get(name) + if rank_slice is not None: + rank_dim, active_rank = rank_slice + tensor = _slice_lora_tensor_to_rank(tensor, rank_dim, active_rank) + lora_state_dict[key] = tensor.cpu() return lora_state_dict @@ -384,12 +518,113 @@ def _convert_from_peft_lora_key(name: str) -> str: return name +def _device_mesh_local_index(mesh) -> int: + """Return this rank's coordinate in a 1D DeviceMesh.""" + try: + return int(mesh.get_local_rank(mesh_dim=0)) + except Exception: + pass + + try: + coordinate = mesh.get_coordinate() + if coordinate is not None: + return int(coordinate[0]) + except Exception: + pass + + try: + return int(dist.get_rank(mesh.get_group())) + except Exception: + return 0 + + +def get_lora_tensor_shard_specs( + model: nn.Module, + names: Optional[Iterable[str]] = None, +) -> Dict[str, LoraTensorShardSpec]: + """Return EP shard specs for LoRA parameters in ``model`` keyed by parameter name.""" + requested_names = set(names) if names is not None else None + fqn2spec_info = getattr(model, "_fqn2spec_info", None) + shard_specs: Dict[str, LoraTensorShardSpec] = {} + + for name, param in model.named_parameters(): + if "lora_A" not in name and "lora_B" not in name: + continue + + clean_name = name.replace("_fsdp_wrapped_module.", "").replace("_orig_mod.", "") + if requested_names is not None and name not in requested_names and clean_name not in requested_names: + continue + + spec_info = None + if fqn2spec_info is not None: + spec_info = fqn2spec_info.get(clean_name) or fqn2spec_info.get(name) + if spec_info is None: + spec_info = getattr(param, "spec_info", None) + if spec_info is None or not isinstance(spec_info.placement, Shard): + continue + + ep_mesh = spec_info.ep_mesh + if ep_mesh is None: + continue + + shard_specs[name] = LoraTensorShardSpec( + dim=spec_info.placement.dim, + index=_device_mesh_local_index(ep_mesh), + size=int(ep_mesh.size()), + ) + if clean_name != name: + shard_specs[clean_name] = shard_specs[name] + + return shard_specs + + +def _slice_lora_tensor_shard( + key: str, + tensor: torch.Tensor, + expected_shape: Tuple[int, ...], + shard_specs: Optional[Dict[str, LoraTensorShardSpec]], +) -> Optional[torch.Tensor]: + if shard_specs is None or key not in shard_specs: + return None + + spec = shard_specs[key] + shard_dim = spec.dim if spec.dim >= 0 else tensor.dim() + spec.dim + if shard_dim < 0 or shard_dim >= tensor.dim(): + return None + if len(expected_shape) != tensor.dim(): + return None + if spec.size <= 1: + return None + + for dim, size in enumerate(expected_shape): + if dim != shard_dim and tensor.shape[dim] != size: + return None + + total_size = tensor.shape[shard_dim] + shard_size = (total_size + spec.size - 1) // spec.size + start = spec.index * shard_size + if start >= total_size: + return None + + length = min(shard_size, total_size - start) + sliced = tensor.narrow(shard_dim, start, length).contiguous() + if tuple(sliced.shape) == expected_shape: + return sliced + + raise RuntimeError( + f"Converted LoRA tensor shard for {key} has shape {tuple(sliced.shape)}, " + f"expected {expected_shape} (full shape {tuple(tensor.shape)}, " + f"shard dim {shard_dim}, shard {spec.index}/{spec.size})" + ) + + def _align_lora_tensor_shape( key: str, tensor: torch.Tensor, expected_shapes: Optional[Dict[str, torch.Size]], + expected_shard_specs: Optional[Dict[str, LoraTensorShardSpec]] = None, ) -> torch.Tensor: - """Match a converted tensor to the live LoRA shape, allowing a final-dim transpose.""" + """Match a converted tensor to the live LoRA shape, allowing transpose and EP slicing.""" if expected_shapes is None or key not in expected_shapes: return tensor @@ -397,17 +632,32 @@ def _align_lora_tensor_shape( if tuple(tensor.shape) == expected_shape: return tensor + sliced = _slice_lora_tensor_shard(key, tensor, expected_shape, expected_shard_specs) + if sliced is not None: + return sliced + if tensor.dim() >= 2: transposed = tensor.transpose(-2, -1).contiguous() if tuple(transposed.shape) == expected_shape: return transposed + sliced = _slice_lora_tensor_shard(key, transposed, expected_shape, expected_shard_specs) + if sliced is not None: + return sliced + + if expected_shard_specs is not None and key in expected_shard_specs: + spec = expected_shard_specs[key] + raise RuntimeError( + f"Converted LoRA tensor for {key} has shape {tuple(tensor.shape)}, expected {expected_shape} " + f"(shard dim {spec.dim}, shard {spec.index}/{spec.size})" + ) 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, + expected_shard_specs: Optional[Dict[str, LoraTensorShardSpec]] = None, ) -> Dict[str, torch.Tensor]: """Convert PEFT-style LoRA checkpoint tensors into xorl's internal layout. @@ -422,6 +672,9 @@ def convert_peft_lora_state_dict( 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. + expected_shard_specs: Optional EP shard specs keyed by internal name. + When provided, full MoE expert tensors are sliced to the current + rank's local expert shard after any required PEFT transpose. Returns: State dict keyed by xorl's internal LoRA parameter names. @@ -435,6 +688,19 @@ def convert_peft_lora_state_dict( else: key = raw_key + sglang_match = _MOE_SGLANG_SHARED_OUTER_PATTERN.match(key) + if sglang_match is not None: + prefix, w_slot, lora_type = sglang_match.groups() + proj_name = _SGLANG_W_TO_PROJ[w_slot] + internal_name = f"{prefix}.mlp.experts.{proj_name}_lora_{lora_type}" + # shared_outer stores 3D tensors transposed (last two dims) vs. + # xorl's in-memory layout. Flip them back to the in-first order. + restored = value.transpose(-2, -1).contiguous() if value.dim() >= 2 else value + converted_state_dict[internal_name] = _align_lora_tensor_shape( + internal_name, restored, expected_shapes, expected_shard_specs + ) + continue + match = _MOE_PEFT_LORA_PATTERN.match(key) if match is not None: prefix, expert_token, proj_name, lora_type = match.groups() @@ -446,7 +712,9 @@ def convert_peft_lora_state_dict( continue internal_name = _convert_from_peft_lora_key(key) - converted_state_dict[internal_name] = _align_lora_tensor_shape(internal_name, value, expected_shapes) + converted_state_dict[internal_name] = _align_lora_tensor_shape( + internal_name, value, expected_shapes, expected_shard_specs + ) for internal_name, bucket in moe_buckets.items(): if "shared" in bucket: @@ -464,7 +732,9 @@ def convert_peft_lora_state_dict( ) 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) + converted_state_dict[internal_name] = _align_lora_tensor_shape( + internal_name, stacked, expected_shapes, expected_shard_specs + ) return converted_state_dict @@ -478,7 +748,9 @@ 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, + transpose_moe_lora_to_peft: bool = True, + lora_export_format: str = "peft", + preserve_lora_dtype: bool = False, ) -> str: """ Save LoRA weights in PEFT-compatible format. @@ -499,18 +771,43 @@ def save_lora_checkpoint( 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. + into PEFT/SGLang orientation during export. Enabled by default so + exported MoE adapters match the shape convention expected by + inference backends. Ignored when + ``lora_export_format="sglang_shared_outer"``. + lora_export_format: On-disk layout for MoE expert LoRA. ``"peft"`` + (default) un-stacks the 3D tensors into per-expert 2D keys. Pass + ``"sglang_shared_outer"`` to emit SGLang's stacked 3D shared_outer + layout directly (requires ``moe_hybrid_shared_lora=True``). + preserve_lora_dtype: Keep LoRA tensor dtypes in the safetensors file + instead of exporting bf16 weights. Use this for training-resume + checkpoints; keep the default bf16 export for inference adapters. Returns: Path to saved checkpoint directory """ + if lora_export_format not in LORA_EXPORT_FORMATS: + raise ValueError(f"Unknown lora_export_format={lora_export_format!r}. Expected one of {LORA_EXPORT_FORMATS}.") + if lora_export_format == "sglang_shared_outer" and not moe_hybrid_shared_lora: + raise ValueError( + "lora_export_format='sglang_shared_outer' requires moe_hybrid_shared_lora=True " + "(shared_outer only makes sense for hybrid-shared MoE LoRA)." + ) + os.makedirs(save_path, exist_ok=True) # Get LoRA state dict — use provided one or extract from model if lora_state_dict is None: lora_state_dict = get_lora_state_dict(model) + else: + lora_state_dict = slice_lora_state_dict_to_active_rank(model, lora_state_dict) + + def _prepare_lora_tensor(tensor: torch.Tensor) -> torch.Tensor: + tensor = tensor.detach().cpu().contiguous() + if preserve_lora_dtype: + return tensor + return tensor.to(torch.bfloat16) def _is_moe_lora_param(name: str) -> bool: """Check if this is a stacked MoE LoRA parameter.""" @@ -568,7 +865,7 @@ def _unmerge_moe_lora_weights(name: str, stacked_tensor: torch.Tensor) -> Dict[s 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) + result[peft_key] = _prepare_lora_tensor(expert_tensor) else: # Per-expert weights: use expert index in the key name for expert_idx in range(num_experts): @@ -578,10 +875,25 @@ def _unmerge_moe_lora_weights(name: str, stacked_tensor: torch.Tensor) -> Dict[s # 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" - result[peft_key] = expert_tensor.to(torch.bfloat16) + result[peft_key] = _prepare_lora_tensor(expert_tensor) return result + def _sglang_shared_outer_moe_weight(name: str, stacked_tensor: torch.Tensor) -> Tuple[str, torch.Tensor]: + """Rename + transpose a stacked MoE tensor into SGLang shared_outer layout. + + xorl stores 3D tensors in-first (A: [E_or_1, in, r], B: [E_or_1, r, out]). + shared_outer stores them out-first under ``experts.w{1,2,3}.lora_{A|B}.weight``. + """ + match = _MOE_LORA_PATTERN.match(name) + if not match: + raise ValueError(f"Invalid MoE LoRA parameter name: {name}") + prefix, proj_name, lora_type = match.group(1), match.group(2), match.group(3) + w_slot = _PROJ_TO_SGLANG_W[proj_name] + peft_key = f"{_PEFT_BASE_MODEL_PREFIX}{prefix}.mlp.experts.{w_slot}.lora_{lora_type}.weight" + out_tensor = _prepare_lora_tensor(stacked_tensor.transpose(-2, -1).contiguous()) + return peft_key, out_tensor + # Convert keys to PEFT format: base_model.model.{converted_key} peft_state_dict = {} detected_modules = set() @@ -590,18 +902,20 @@ def _unmerge_moe_lora_weights(name: str, stacked_tensor: torch.Tensor) -> Dict[s for key, value in lora_state_dict.items(): # Check if this is a stacked MoE LoRA parameter if _is_moe_lora_param(key): - # Unmerge stacked MoE LoRA weights into per-expert format - per_expert_weights = _unmerge_moe_lora_weights(key, value) - peft_state_dict.update(per_expert_weights) + if lora_export_format == "sglang_shared_outer": + peft_key, out_tensor = _sglang_shared_outer_moe_weight(key, value) + peft_state_dict[peft_key] = out_tensor + else: + # Unmerge stacked MoE LoRA weights into per-expert format + 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) 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": # 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] + detected_r = value.shape[2] else: # Extract module name for target_modules detection parts = key.split(".") @@ -616,18 +930,32 @@ def _unmerge_moe_lora_weights(name: str, stacked_tensor: torch.Tensor) -> Dict[s # Convert to PEFT key format peft_key = f"base_model.model.{_convert_to_peft_key(key)}" - peft_state_dict[peft_key] = value.to(torch.bfloat16) + peft_state_dict[peft_key] = _prepare_lora_tensor(value) # Auto-detect parameters if not provided if target_modules is None: target_modules = list(detected_modules) - if r is None: - r = detected_r or 16 - if lora_alpha is None: + active_rank, active_alpha = _first_active_lora_config(model) + if detected_r is not None: + if r is not None and r != detected_r: + logger.warning( + f"Requested LoRA config r={r} does not match exported tensor rank {detected_r}; writing r={detected_r}." + ) + r = detected_r + elif r is None: + r = active_rank or 16 + if active_alpha is not None and active_rank == r: + if lora_alpha is not None and lora_alpha != active_alpha: + logger.warning( + f"Requested LoRA alpha={lora_alpha} does not match active alpha {active_alpha}; " + f"writing lora_alpha={active_alpha}." + ) + lora_alpha = active_alpha + elif lora_alpha is None: # Try to detect from model for module in model.modules(): if isinstance(module, LoraLinear): - lora_alpha = module.lora_alpha + lora_alpha = getattr(module, "active_lora_alpha", module.lora_alpha) break if lora_alpha is None: lora_alpha = r # Default: alpha = r @@ -649,8 +977,16 @@ def _unmerge_moe_lora_weights(name: str, stacked_tensor: torch.Tensor) -> Dict[s "peft_type": "LORA", "inference_mode": True, "fan_in_fan_out": False, - "moe_hybrid_shared_lora": moe_hybrid_shared_lora, } + if lora_export_format == "sglang_shared_outer": + adapter_config["_sglang_lora_format"] = "shared_outer" + # SGLang's lora_manager classifies adapters via the moe_hybrid_shared_lora + # / shared_moe_lora keys in hf_config. shared_outer IS hybrid_shared + # on-disk, so mirror the flag here so SGLang doesn't mis-classify it as + # per_expert (which would reject loading under --lora-moe-format hybrid_shared). + adapter_config["moe_hybrid_shared_lora"] = True + else: + adapter_config["moe_hybrid_shared_lora"] = moe_hybrid_shared_lora config_path = os.path.join(save_path, "adapter_config.json") with open(config_path, "w") as f: @@ -697,7 +1033,12 @@ def load_lora_checkpoint( 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) + lora_shard_specs = get_lora_tensor_shard_specs(model, names=model_lora_shapes.keys()) + converted_state_dict = convert_peft_lora_state_dict( + state_dict, + expected_shapes=model_lora_shapes, + expected_shard_specs=lora_shard_specs, + ) # Load into model missing, unexpected = load_lora_state_dict(model, converted_state_dict, strict=strict) @@ -828,9 +1169,10 @@ def inject_lora_into_model_with_moe( model: Model to inject LoRA into r: LoRA rank lora_alpha: LoRA alpha for scaling - 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"] + target_modules: List of module names to target. If None, falls back to + the architecture default in DEFAULT_TARGET_MODULES (e.g. + q_a_proj/q_b_proj/kv_a_proj_with_mqa/kv_b_proj/o_proj + for DeepSeek-V3 / Kimi). moe_hybrid_shared_lora: If True, use hybrid sharing (lora_A shared for gate/up, lora_B shared for down) Returns: @@ -846,11 +1188,17 @@ def inject_lora_into_model_with_moe( ... ) """ if target_modules is None: - target_modules = ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"] - - # Separate attention modules from MLP/expert modules - attention_modules = [m for m in target_modules if m in ["q_proj", "k_proj", "v_proj", "o_proj", "lm_head"]] - expert_modules = [m for m in target_modules if m in ["gate_proj", "up_proj", "down_proj"]] + target_modules = _get_default_target_modules(model) + + # Partition target_modules into expert MLP projections (handled by the + # group-GEMM MoE path) and dense linears (attention plus per-layer dense + # MLPs, handled by inject_lora_into_model). Anything that isn't an expert + # name is treated as a dense target so this function inherits new + # architectures (DeepSeek-V3 / Kimi MLA: q_a_proj, q_b_proj, + # kv_a_proj_with_mqa, kv_b_proj) from DEFAULT_TARGET_MODULES rather than + # silently dropping them via a Llama/Qwen-shaped allowlist. + expert_modules = [m for m in target_modules if m in ("gate_proj", "up_proj", "down_proj")] + attention_modules = [m for m in target_modules if m not in expert_modules] # Step 1: Inject LoRA into standard nn.Linear layers # This includes attention projections AND dense MLP layers (for layers without MoE) @@ -890,6 +1238,7 @@ def get_moe_lora_state_dict(model: nn.Module) -> Dict[str, torch.Tensor]: State dict containing MoE LoRA parameters """ moe_lora_state_dict = {} + rank_slices = _active_lora_rank_slices(model) for name, module in model.named_modules(): # Check if this is an MoE LoRA expert module (has lora_config attribute) @@ -898,7 +1247,12 @@ def get_moe_lora_state_dict(model: nn.Module) -> Dict[str, torch.Tensor]: for param_name, param in module.named_parameters(): if "lora_" in param_name: full_key = f"{name}.{param_name}" - moe_lora_state_dict[full_key] = param.detach().cpu() + tensor = param.detach() + rank_slice = rank_slices.get(full_key) + if rank_slice is not None: + rank_dim, active_rank = rank_slice + tensor = _slice_lora_tensor_to_rank(tensor, rank_dim, active_rank) + moe_lora_state_dict[full_key] = tensor.cpu() return moe_lora_state_dict diff --git a/src/xorl/models/__init__.py b/src/xorl/models/__init__.py index 9262632d..7149fb57 100644 --- a/src/xorl/models/__init__.py +++ b/src/xorl/models/__init__.py @@ -2,10 +2,12 @@ from .auto import build_foundation_model, build_processor, build_tokenizer from .module_utils import ( all_ranks_load_weights, + grouped_load_weights, init_empty_weights, rank0_load_and_broadcast_weights, save_model_assets, save_model_weights, + save_model_weights_distributed, ) @@ -15,8 +17,10 @@ "build_tokenizer", "init_empty_weights", "all_ranks_load_weights", + "grouped_load_weights", "rank0_load_and_broadcast_weights", "save_model_assets", + "save_model_weights_distributed", "save_model_weights", "transformers", ] diff --git a/src/xorl/models/auto.py b/src/xorl/models/auto.py index 2015aad4..0d4d8c77 100644 --- a/src/xorl/models/auto.py +++ b/src/xorl/models/auto.py @@ -1,4 +1,6 @@ +import json import types +from pathlib import Path from typing import TYPE_CHECKING, Any, Dict, Literal, Optional, Union import torch @@ -16,6 +18,10 @@ from .layers.attention import ATTENTION_FUNCTIONS from .layers.normalization import set_rmsnorm_mode from .loader import ModelLoader, get_loader +from .transformers.deepseek_v3.configuration_deepseek_v3 import DeepseekV3Config +from .transformers.deepseek_v3.support import validate_deepseek_v3_router_settings +from .transformers.glm4_moe.configuration_glm4_moe import Glm4MoeConfig +from .transformers.gpt_oss.configuration_gpt_oss import GptOssConfig 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 ( @@ -30,6 +36,33 @@ logger = logging.get_logger(__name__) +def _build_local_kimi_tokenizer(tokenizer_path: str): + tokenizer_dir = Path(tokenizer_path) + tokenizer_config_path = tokenizer_dir / "tokenizer_config.json" + vocab_file = tokenizer_dir / "tiktoken.model" + if not tokenizer_config_path.is_file() or not vocab_file.is_file(): + return None + + with tokenizer_config_path.open() as f: + tokenizer_config = json.load(f) + + auto_tokenizer = tokenizer_config.get("auto_map", {}).get("AutoTokenizer", []) + auto_tokenizer_cls = auto_tokenizer[0] if auto_tokenizer else "" + if tokenizer_config.get("tokenizer_class") != "TikTokenTokenizer" and not auto_tokenizer_cls.endswith( + "TikTokenTokenizer" + ): + return None + + from .transformers.deepseek_v3.tokenization_kimi import TikTokenTokenizer # noqa: PLC0415 + + tokenizer_kwargs = dict(tokenizer_config) + tokenizer_kwargs.pop("auto_map", None) + tokenizer_kwargs.pop("tokenizer_class", None) + tokenizer_kwargs.pop("vocab_file", None) + tokenizer_kwargs["padding_side"] = "right" + return TikTokenTokenizer(vocab_file=str(vocab_file), **tokenizer_kwargs) + + def _namespace_from_dict(value): if isinstance(value, dict): return types.SimpleNamespace(**{k: _namespace_from_dict(v) for k, v in value.items()}) @@ -45,24 +78,53 @@ def _load_local_xorl_config( config_dict, _ = PretrainedConfig.get_config_dict(config_path, **config_kwargs) model_type = config_dict.get("model_type") + if model_type == "glm4_moe": + return Glm4MoeConfig.from_dict(config_dict) + if model_type == "qwen3_5_moe": return Qwen3_5MoeConfig.from_hf_config(_namespace_from_dict(config_dict)) if model_type == "qwen3_5": return Qwen3_5Config.from_hf_config(_namespace_from_dict(config_dict)) + if model_type in {"deepseek_v3", "kimi_k2", "kimi_k25"}: + return DeepseekV3Config.from_hf_config(_namespace_from_dict(config_dict)) + if model_type == "qwen2": - from .transformers.qwen2.configuration_qwen2 import Qwen2Config + from .transformers.qwen2.configuration_qwen2 import Qwen2Config # noqa: PLC0415 return Qwen2Config(**{k: v for k, v in config_dict.items() if not k.startswith("_")}) + if model_type == "gpt_oss": + return GptOssConfig.from_hf_config(_namespace_from_dict(config_dict)) + if model_type == "olmo2": + from .transformers.olmo2.configuration_olmo2 import Olmo2Config # noqa: PLC0415 + + return Olmo2Config(**{k: v for k, v in config_dict.items() if not k.startswith("_")}) + return None +def _get_architectures(config: "PretrainedConfig") -> set[str]: + architectures = getattr(config, "architectures", None) + if architectures is None: + return set() + if isinstance(architectures, list): + return set(architectures) + return {architectures} + + +def _is_gpt_oss_config(config: "PretrainedConfig") -> bool: + return getattr(config, "model_type", None) == "gpt_oss" or "GptOssForCausalLM" in _get_architectures(config) + + def build_tokenizer(tokenizer_path: str) -> "PreTrainedTokenizer": """ Builds the tokenizer. """ + tokenizer = _build_local_kimi_tokenizer(tokenizer_path) + if tokenizer is not None: + return tokenizer return AutoTokenizer.from_pretrained(tokenizer_path, padding_side="right") @@ -110,6 +172,7 @@ def build_foundation_model( moe_implementation: Optional[Literal["eager", "triton", "native", "quack"]] = None, ep_dispatch: str = "alltoall", train_router: bool = False, + record_routing_weights: bool = True, deepep_buffer_size_gb: float = 2.0, deepep_num_sms: int = 20, deepep_async_combine: bool = False, @@ -119,6 +182,7 @@ def build_foundation_model( activation_native: bool = False, rope_native: bool = False, attention_cast_bf16: bool = False, + flash_attention_deterministic: bool = False, init_device: Literal["cpu", "cuda", "npu", "meta"] = "cuda", config_kwargs: Optional[Dict[str, Any]] = None, ) -> nn.Module: @@ -143,6 +207,8 @@ def build_foundation_model( config._moe_implementation = moe_implementation logger.info_rank0(f"Moe implementation: {moe_implementation}") + validate_deepseek_v3_router_settings(config, train_router=train_router) + if ep_dispatch == "deepep" and train_router: raise ValueError( "train_router=True is not supported with ep_dispatch='deepep'. " @@ -151,6 +217,7 @@ def build_foundation_model( config._ep_dispatch = ep_dispatch config.train_router = train_router + config.record_routing_weights = record_routing_weights config._deepep_buffer_size_gb = deepep_buffer_size_gb config._deepep_num_sms = deepep_num_sms config._deepep_async_combine = deepep_async_combine @@ -161,6 +228,7 @@ def build_foundation_model( config._activation_native = activation_native config._rope_native = rope_native config._attention_cast_bf16 = attention_cast_bf16 + config._flash_attention_deterministic = flash_attention_deterministic if ep_dispatch == "deepep": logger.info_rank0( @@ -180,6 +248,14 @@ def build_foundation_model( logger.warning_once(LINEAR_ATTENTION_RING_UNSUPPORTED_MESSAGE) raise ValueError(LINEAR_ATTENTION_RING_UNSUPPORTED_MESSAGE) + if _is_gpt_oss_config(config) and attn_implementation not in ("eager", "flash_attention_3"): + raise ValueError( + "GPT-OSS attention sinks are only implemented for attn_implementation=" + "'eager' or 'flash_attention_3' in xorl. Using other backends (sdpa, " + "flash_attention_2, flash_attention_4, native) would silently drop the " + "sink logits and change model outputs." + ) + loader: ModelLoader = get_loader(config) # Validate FA4 availability early diff --git a/src/xorl/models/base.py b/src/xorl/models/base.py index 2d5b76a5..ae3db928 100644 --- a/src/xorl/models/base.py +++ b/src/xorl/models/base.py @@ -8,6 +8,7 @@ from ..utils import logging from .layers.moe.moe_block import MoEBlock from .layers.moe.routing_replay import RoutingReplay +from .module_utils import DEFAULT_GRADIENT_CHECKPOINTING_METHOD, GradientCheckpointingMethod logger = logging.get_logger(__name__) @@ -36,6 +37,7 @@ def __init__(self, config): super().__init__() self.config = config self.gradient_checkpointing = False + self._gradient_checkpointing_method: GradientCheckpointingMethod | None = None # ------------------------------------------------------------------ # Weight initialisation @@ -91,8 +93,13 @@ def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs=None): if gradient_checkpointing_kwargs is None: gradient_checkpointing_kwargs = {"use_reentrant": False} - # Pop selective checkpoint config before passing to torch checkpoint - grad_ckpt_method = gradient_checkpointing_kwargs.pop("gradient_checkpointing_method", None) + # Pop selective checkpoint config before passing to torch checkpoint. + # The default must be a valid method string — writing `None` here + # poisons every gate's equality check and silently disables + # checkpointing (see scratch/rl/debug_log/xorl-rebase-root-cause.md). + grad_ckpt_method = gradient_checkpointing_kwargs.pop( + "gradient_checkpointing_method", DEFAULT_GRADIENT_CHECKPOINTING_METHOD + ) grad_ckpt_func = partial(torch.utils.checkpoint.checkpoint, **gradient_checkpointing_kwargs) @@ -102,7 +109,7 @@ def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs=None): module._gradient_checkpointing_func = grad_ckpt_func module._gradient_checkpointing_method = grad_ckpt_method - if grad_ckpt_method is not None: + if grad_ckpt_method != DEFAULT_GRADIENT_CHECKPOINTING_METHOD: logger.info(f"Selective checkpointing enabled: gradient_checkpointing_method={grad_ckpt_method}") # Create routing replay instances for MoE blocks diff --git a/src/xorl/models/checkpoint_handlers/buffers.py b/src/xorl/models/checkpoint_handlers/buffers.py index 58d562eb..e0fb61c7 100644 --- a/src/xorl/models/checkpoint_handlers/buffers.py +++ b/src/xorl/models/checkpoint_handlers/buffers.py @@ -382,11 +382,20 @@ def add(self, layer_idx: int, expert_idx: int, proj: str, tensor: torch.Tensor) if expert_idx < self.expert_start or expert_idx >= self.expert_end: return - # GPU fast path: pin+DMA to GPU, transpose on GPU (~13x faster than CPU) + # GPU fast path: transpose and stack on GPU. If the tensor is already on + # the target GPU (for example after inline dequantization), keep it there + # instead of bouncing through CPU memory again. if self._device is not None and self._device.type == "cuda": - if tensor.dtype != torch.bfloat16: - tensor = tensor.to(dtype=torch.bfloat16) - tensor = tensor.pin_memory().to(device=self._device, non_blocking=True) + target_dtype = tensor.dtype if tensor.is_floating_point() else torch.bfloat16 + if tensor.device == self._device: + if tensor.dtype != target_dtype: + tensor = tensor.to(dtype=target_dtype) + elif tensor.device.type == "cpu": + if tensor.dtype != target_dtype: + tensor = tensor.to(dtype=target_dtype) + tensor = tensor.pin_memory().to(device=self._device, non_blocking=True) + else: + tensor = tensor.to(device=self._device, dtype=target_dtype, non_blocking=True) tensor = tensor.t().contiguous() if key not in self._stacked_buffers: diff --git a/src/xorl/models/layers/attention/backend/flash_attention.py b/src/xorl/models/layers/attention/backend/flash_attention.py index 7de023e7..ccd57800 100644 --- a/src/xorl/models/layers/attention/backend/flash_attention.py +++ b/src/xorl/models/layers/attention/backend/flash_attention.py @@ -27,6 +27,7 @@ # Environment variable to disable FA4 even when available XORL_DISABLE_FA4 = os.environ.get("XORL_DISABLE_FA4", "0") == "1" +XORL_FLASH_ATTN_DETERMINISTIC = os.environ.get("XORL_FLASH_ATTN_DETERMINISTIC", "0") == "1" logger = logging.get_logger(__name__) @@ -66,6 +67,13 @@ def flash_attention_forward( position_ids = kwargs.pop("position_ids", None) if position_ids is not None and position_ids.dim() == 3: position_ids = position_ids[0] + deterministic = bool( + kwargs.pop( + "deterministic", + XORL_FLASH_ATTN_DETERMINISTIC + or getattr(getattr(module, "config", None), "_flash_attention_deterministic", False), + ) + ) # FA4 (CUTE) path if _should_use_fa4(use_fa4): @@ -160,6 +168,7 @@ def flash_attention_forward( causal=causal, window_size=window_size_fa3, softcap=softcap if softcap is not None else 0.0, + deterministic=deterministic, ) # Restore batch dimension attn_output = attn_output.unsqueeze(0) @@ -174,6 +183,7 @@ def flash_attention_forward( causal=causal, window_size=window_size_fa3, softcap=softcap if softcap is not None else 0.0, + deterministic=deterministic, ) return attn_output, None diff --git a/src/xorl/models/layers/moe/aux_loss.py b/src/xorl/models/layers/moe/aux_loss.py index 0cce91fe..ec450f02 100644 --- a/src/xorl/models/layers/moe/aux_loss.py +++ b/src/xorl/models/layers/moe/aux_loss.py @@ -48,6 +48,10 @@ def global_load_balancing_loss_func( if gate_logits is None or not isinstance(gate_logits, tuple): return 0 + gate_logits = tuple(layer_gate for layer_gate in gate_logits if layer_gate is not None) + if not gate_logits: + 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) diff --git a/src/xorl/models/layers/moe/backend/__init__.py b/src/xorl/models/layers/moe/backend/__init__.py index 5b8f0796..51dc7ba9 100644 --- a/src/xorl/models/layers/moe/backend/__init__.py +++ b/src/xorl/models/layers/moe/backend/__init__.py @@ -52,31 +52,80 @@ try: from xorl.ops.moe.triton import TritonEPGroupGemm - EP_EXPERT_COMPUTE["triton"] = TritonEPGroupGemm.apply + def _triton_ep_apply( + permute_tokens, + cumsum, + gate_up_proj, + down_proj, + intermediate_size, + expert_scores=None, + hidden_act="silu", + **_extras, # GPT-OSS extras (gate_up_bias, down_bias) — triton doesn't support them + ): + if any(v is not None for v in _extras.values()) or hidden_act == "clamped_swiglu": + raise NotImplementedError( + "triton EP backend does not support per-expert biases or clamped_swiglu " + "activation (required by GPT-OSS). Use moe_implementation='native' instead." + ) + return TritonEPGroupGemm.apply( + permute_tokens, cumsum, gate_up_proj, down_proj, intermediate_size, expert_scores, hidden_act + ) + + EP_EXPERT_COMPUTE["triton"] = _triton_ep_apply except ImportError: pass -# Quack EP compute — adapt fused interface to old (gate_proj, up_proj) signature +# Quack EP compute — fused gate_up_proj interface try: 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) + def _quack_ep_fused( + permute_tokens, + cumsum, + gate_up_proj, + down_proj, + intermediate_size, + expert_scores=None, + hidden_act="silu", + **_extras, # GPT-OSS extras (gate_up_bias, down_bias) — quack doesn't support them + ): + if any(v is not None for v in _extras.values()) or hidden_act == "clamped_swiglu": + raise NotImplementedError( + "quack EP backend does not support per-expert biases or clamped_swiglu " + "activation (required by GPT-OSS). Use moe_implementation='native' instead." + ) + return _QuackEPGroupGemm.apply( + permute_tokens, cumsum, gate_up_proj, down_proj, intermediate_size, expert_scores, hidden_act + ) EP_EXPERT_COMPUTE["quack"] = _quack_ep_fused except ImportError: pass -# Native EP compute — adapt fused interface +# Native EP compute — fused gate_up_proj interface try: from .native import native_ep_compute as _native_ep_compute - def _native_ep_fused(permute_tokens, cumsum, gate_up_proj, down_proj, intermediate_size, 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) + def _native_ep_fused( + permute_tokens, + cumsum, + gate_up_proj, + down_proj, + intermediate_size, + expert_scores=None, + hidden_act="silu", + **extras, # gate_up_bias, down_bias — optional, used for GPT-OSS + ): + del intermediate_size + return _native_ep_compute( + permute_tokens, + cumsum, + gate_up_proj, + down_proj, + expert_scores, + hidden_act=hidden_act, + **extras, + ) EP_EXPERT_COMPUTE["native"] = _native_ep_fused except ImportError: diff --git a/src/xorl/models/layers/moe/backend/eager.py b/src/xorl/models/layers/moe/backend/eager.py index 6cf0a0cc..cc98de0b 100644 --- a/src/xorl/models/layers/moe/backend/eager.py +++ b/src/xorl/models/layers/moe/backend/eager.py @@ -3,6 +3,7 @@ import torch from xorl.distributed.parallel_state import get_parallel_state +from xorl.ops.moe.activations import apply_moe_activation def eager_expert_forward( @@ -11,7 +12,9 @@ def eager_expert_forward( gate_proj: torch.Tensor, up_proj: torch.Tensor, down_proj: torch.Tensor, - act_fn, + hidden_act: str = "silu", + gate_up_bias: torch.Tensor | None = None, + down_bias: torch.Tensor | None = None, ) -> torch.Tensor: """Forward pass for a single expert (eager mode). @@ -23,14 +26,21 @@ def eager_expert_forward( gate_proj: Gate projection weights ``[num_experts, hidden, intermediate]``. up_proj: Up projection weights ``[num_experts, hidden, intermediate]``. down_proj: Down projection weights ``[num_experts, intermediate, hidden]``. - act_fn: Activation function (e.g. ``torch.nn.SiLU``). - - Returns: - Output tensor of shape ``(num_tokens, hidden_dim)``. + hidden_act: Activation kind (e.g. ``"silu"``, ``"gelu_tanh"``, + ``"clamped_swiglu"``) — dispatched via ``apply_moe_activation``. + gate_up_bias: Optional per-expert bias ``[num_experts, 2*intermediate]``, + split as ``[gate_bias | up_bias]`` along the last dim. + down_bias: Optional per-expert bias ``[num_experts, hidden_dim]``. """ 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 + if gate_up_bias is not None: + intermediate_size = gate_proj_out.shape[-1] + gate_proj_out = gate_proj_out + gate_up_bias[expert_idx, :intermediate_size] + up_proj_out = up_proj_out + gate_up_bias[expert_idx, intermediate_size:] + out = apply_moe_activation(hidden_act, gate_proj_out, up_proj_out) out = torch.matmul(out, down_proj[expert_idx]) + if down_bias is not None: + out = out + down_bias[expert_idx] return out diff --git a/src/xorl/models/layers/moe/backend/native.py b/src/xorl/models/layers/moe/backend/native.py index 118fb395..3ad399ab 100644 --- a/src/xorl/models/layers/moe/backend/native.py +++ b/src/xorl/models/layers/moe/backend/native.py @@ -128,40 +128,61 @@ def _run_experts_grouped_mm( x: torch.Tensor, padded_counts: torch.Tensor, expert_scores: torch.Tensor | None = None, + hidden_act: str = "silu", + gate_up_proj: torch.Tensor | None = None, + gate_up_bias: torch.Tensor | None = None, + down_bias: torch.Tensor | None = None, ) -> torch.Tensor: """Run MoE experts using ``torch._grouped_mm``. Compiled with ``torch.compile(fullgraph=True)`` for operator fusion. - Weight shapes in (G, K, N) format:: + When ``gate_up_proj`` is provided, uses a single fused GEMM (matching HF) + instead of two separate gate/up GEMMs for better bf16 numerical consistency. - gate_proj: [num_experts, hidden_dim, intermediate_size] - up_proj: [num_experts, hidden_dim, intermediate_size] - down_proj: [num_experts, intermediate_size, hidden_dim] + Optional per-expert biases (pre-expanded to ``[total_padded, dim]``) are + applied before the activation (gate_up_bias) and after the down projection + (down_bias). """ offsets = torch.cumsum(padded_counts, dim=0, dtype=torch.int32) compute_dtype = torch.bfloat16 - # gate: x @ gate_proj -> (tokens, intermediate) - gate_out = F.silu( - torch._grouped_mm( + if gate_up_proj is not None: + # Fused: single GEMM -> chunk (matches HF's grouped_mm dispatch) + gate_up_out = torch._grouped_mm( x.to(compute_dtype), - gate_proj.to(compute_dtype), + gate_up_proj.to(compute_dtype), offs=offsets, ) - ) - - # up: x @ up_proj -> (tokens, intermediate) - up_out = torch._grouped_mm( - x.to(compute_dtype), - up_proj.to(compute_dtype), - offs=offsets, - ) - - # SwiGLU: silu(gate) * up - h = gate_out * up_out - if expert_scores is not None: - h = h * expert_scores.to(h.dtype).unsqueeze(-1) + intermediate_size = gate_up_out.shape[-1] // 2 + gate_raw = gate_up_out[..., :intermediate_size] + up_out = gate_up_out[..., intermediate_size:] + else: + # Split: separate gate/up GEMMs (legacy path) + gate_raw = torch._grouped_mm(x.to(compute_dtype), gate_proj.to(compute_dtype), offs=offsets) + up_out = torch._grouped_mm(x.to(compute_dtype), up_proj.to(compute_dtype), offs=offsets) + + if gate_up_bias is not None: + intermediate = gate_raw.shape[-1] + gate_raw = gate_raw + gate_up_bias[:, :intermediate].to(compute_dtype) + up_out = up_out + gate_up_bias[:, intermediate:].to(compute_dtype) + + # GLU: act(gate) * up + if hidden_act == "gelu_tanh": + gate_out = F.gelu(gate_raw, approximate="tanh") + h = gate_out * up_out + elif hidden_act == "clamped_swiglu": + # GPT-OSS clamped SwiGLU: clamp both branches, scaled sigmoid gate, + # +1 bias on the up/linear branch. + _CLAMPED_SWIGLU_ALPHA = 1.702 + _CLAMPED_SWIGLU_LIMIT = 7.0 + gate_raw = gate_raw.clamp(max=_CLAMPED_SWIGLU_LIMIT) + up_out = up_out.clamp(min=-_CLAMPED_SWIGLU_LIMIT, max=_CLAMPED_SWIGLU_LIMIT) + gate_out = gate_raw * torch.sigmoid(_CLAMPED_SWIGLU_ALPHA * gate_raw) + h = gate_out * (up_out + 1) + else: + gate_out = F.silu(gate_raw) + h = gate_out * up_out # down: h @ down_proj -> (tokens, hidden) out = torch._grouped_mm( @@ -170,13 +191,43 @@ def _run_experts_grouped_mm( offs=offsets, ).to(x.dtype) + if down_bias is not None: + out = out + down_bias.to(out.dtype) + + # Routing weight must be applied AFTER down_proj + down_bias, not before + # the bias add — otherwise down_bias contributes the same magnitude + # regardless of routing weight (matters when down_bias is non-zero, e.g. + # for GPT-OSS). + if expert_scores is not None: + out = out * expert_scores.to(out.dtype).unsqueeze(-1) + return out -# Compile the inner GEMM function (like torchtitan). +# Compile variants (torch.compile needs static graph). +# Note: torch.compile caches by tensor shapes/strides. The gate_up_proj +# branch is traced correctly on first call with a non-None gate_up_proj. _run_experts_compiled = torch.compile(_run_experts_grouped_mm, fullgraph=True) +def _run_experts_gelu_tanh_wrapper( + gate_proj, up_proj, down_proj, x, padded_counts, expert_scores=None, gate_up_proj=None +): + return _run_experts_grouped_mm( + gate_proj, + up_proj, + down_proj, + x, + padded_counts, + expert_scores, + hidden_act="gelu_tanh", + gate_up_proj=gate_up_proj, + ) + + +_run_experts_compiled_gelu_tanh = torch.compile(_run_experts_gelu_tanh_wrapper, fullgraph=True) + + # --------------------------------------------------------------------------- # Shared token preparation / scatter helper # --------------------------------------------------------------------------- @@ -190,6 +241,15 @@ def _ensure_dynamo_configured(): _native_dynamo_configured = True +@torch.compiler.disable +def _expand_expert_bias( + bias: torch.Tensor, + padded_counts: torch.Tensor, +) -> torch.Tensor: + """Expand per-expert bias ``[E, dim]`` to per-padded-token ``[total_padded, dim]``.""" + return torch.repeat_interleave(bias, padded_counts.long(), dim=0) + + def _native_expert_forward_impl( hidden_states: torch.Tensor, routing_weights: torch.Tensor, @@ -199,12 +259,18 @@ def _native_expert_forward_impl( down_proj: torch.Tensor, num_experts: int, compute_fn, + gate_up_proj: torch.Tensor | None = None, + gate_up_bias: torch.Tensor | None = None, + down_bias: torch.Tensor | None = None, ) -> torch.Tensor: """Shared token-sort / pad / scatter logic for native expert forward. ``compute_fn(gate_proj, up_proj, down_proj, sorted_hidden_padded, padded_counts)`` is called with the prepared padded token tensor and must return a padded output of the same shape. + + Optional ``gate_up_bias`` / ``down_bias`` are per-expert ``[E, dim]`` + tensors expanded to per-padded-token before being passed to ``compute_fn``. """ _ensure_dynamo_configured() num_tokens, top_k = selected_experts.shape @@ -245,9 +311,12 @@ def _native_expert_forward_impl( sorted_hidden_padded = sorted_hidden.new_zeros(total_padded, hidden_dim) 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) + # 7. Expert compute (backend-specific). Pass expert_scores=None — routing + # weights are applied externally below via the deterministic reshape+sum + # path. Applying them twice (inside _run_experts_grouped_mm AND here) would + # square the weights. + gate_up_bias_exp = _expand_expert_bias(gate_up_bias, padded_counts) if gate_up_bias is not None else None + down_bias_exp = _expand_expert_bias(down_bias, padded_counts) if down_bias is not None else None expert_out_padded = compute_fn( gate_proj, @@ -255,15 +324,23 @@ def _native_expert_forward_impl( down_proj, sorted_hidden_padded, padded_counts, - expert_scores_padded, + None, # expert_scores applied after, not inside GEMM + gate_up_proj=gate_up_proj, + gate_up_bias=gate_up_bias_exp, + down_bias=down_bias_exp, ) - # 8. Gather from padded layout (reuse pad_dst) + # 8. Gather from padded layout expert_out = expert_out_padded[pad_dst] - # 9. Scatter-add back to original token positions - output = hidden_states.new_zeros(num_tokens, hidden_dim) - output.index_add_(0, token_ids, expert_out) + # 9. Apply routing weights and accumulate via reshape+sum (deterministic) + expert_out = expert_out * sorted_weights.to(expert_out.dtype).unsqueeze(-1) + + # Unsort back to original (token, top_k_slot) order, then reshape+sum + inv_sorted = torch.empty_like(sorted_order) + inv_sorted[sorted_order] = torch.arange(sorted_order.size(0), device=device) + expert_out = expert_out[inv_sorted] + output = expert_out.view(num_tokens, top_k, hidden_dim).sum(dim=1) return output @@ -273,6 +350,20 @@ def _native_expert_forward_impl( # --------------------------------------------------------------------------- +SUPPORTED_HIDDEN_ACTS = frozenset({"silu", "gelu_tanh", "clamped_swiglu"}) + + +def _select_compiled_fn(hidden_act: str = "silu"): + """Select the right compute variant based on activation kind.""" + if hidden_act == "gelu_tanh": + # Use uncompiled for now — torch.compile fullgraph has issues with + # the gate_up_proj branch when switching between None/non-None. + return lambda *a, **kw: _run_experts_grouped_mm(*a, hidden_act="gelu_tanh", **kw) + if hidden_act == "clamped_swiglu": + return lambda *a, **kw: _run_experts_grouped_mm(*a, hidden_act="clamped_swiglu", **kw) + return _run_experts_compiled + + def native_expert_forward( hidden_states: torch.Tensor, routing_weights: torch.Tensor, @@ -281,9 +372,18 @@ def native_expert_forward( up_proj: torch.Tensor, down_proj: torch.Tensor, num_experts: int, + hidden_act: str = "silu", + gate_up_proj: torch.Tensor = None, **kwargs, ) -> torch.Tensor: - """Forward pass using native PyTorch ``torch._grouped_mm``.""" + """Forward pass using native PyTorch ``torch._grouped_mm``. + + Accepts optional ``gate_up_bias`` / ``down_bias`` via ``**kwargs`` for + GPT-OSS models that use per-expert biases. + """ + from xorl.ops.moe.triton import check_hidden_act_supported + + check_hidden_act_supported(hidden_act, "native", SUPPORTED_HIDDEN_ACTS) return _native_expert_forward_impl( hidden_states, routing_weights, @@ -292,7 +392,10 @@ def native_expert_forward( up_proj, down_proj, num_experts, - _run_experts_compiled, + _select_compiled_fn(hidden_act), + gate_up_proj=gate_up_proj, + gate_up_bias=kwargs.get("gate_up_bias"), + down_bias=kwargs.get("down_bias"), ) @@ -496,34 +599,29 @@ def native_expert_lora_forward( def native_ep_compute( permute_tokens: torch.Tensor, cumsum: torch.Tensor, - gate_proj: torch.Tensor, - up_proj: torch.Tensor, + gate_up_proj: torch.Tensor, down_proj: torch.Tensor, expert_scores: torch.Tensor | None = None, + hidden_act: str = "silu", + gate_up_bias: torch.Tensor | None = None, + down_bias: torch.Tensor | None = None, ) -> torch.Tensor: - """EP expert compute using ``torch._grouped_mm``. + """EP expert compute using ``torch._grouped_mm`` with fused gate+up GEMM. Same interface as ``TritonEPGroupGemm.apply()`` and ``QuackEPGroupGemm.apply()``. - Tokens have already been dispatched via all-to-all; this only handles - the expert MLP computation. - - Args: - permute_tokens: Dispatched tokens ``[total_local_tokens, hidden_dim]``. - cumsum: Cumulative sum of tokens per local expert ``[num_local_experts]``. - gate_proj: ``[num_local_experts, hidden_dim, intermediate_size]``. - up_proj: ``[num_local_experts, hidden_dim, intermediate_size]``. - down_proj: ``[num_local_experts, intermediate_size, hidden_dim]``. - Returns: - Expert outputs ``[total_local_tokens, hidden_dim]``. + Optional ``gate_up_bias`` ``[num_local_experts, 2*I]`` and ``down_bias`` + ``[num_local_experts, H]`` are per-expert biases for GPT-OSS models. """ + from xorl.ops.moe.triton import check_hidden_act_supported + + check_hidden_act_supported(hidden_act, "native", SUPPORTED_HIDDEN_ACTS) if permute_tokens.shape[0] == 0: return permute_tokens - num_local_experts = gate_proj.shape[0] + num_local_experts = gate_up_proj.shape[0] counts = _cumsum_to_counts(cumsum, num_local_experts) - # Pad for alignment and run compiled grouped GEMM padded_tokens, padded_counts = _pad_to_alignment(permute_tokens, counts, num_local_experts) expert_scores_padded = None if expert_scores is not None: @@ -533,11 +631,24 @@ def native_ep_compute( 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 + + gate_up_bias_exp = _expand_expert_bias(gate_up_bias, padded_counts) if gate_up_bias is not None else None + down_bias_exp = _expand_expert_bias(down_bias, padded_counts) if down_bias is not None else None + + compiled_fn = _select_compiled_fn(hidden_act) + # gate_proj/up_proj positional args are unused when gate_up_proj is provided + out_padded = compiled_fn( + None, + None, + down_proj, + padded_tokens, + padded_counts, + expert_scores_padded, + gate_up_proj=gate_up_proj, + gate_up_bias=gate_up_bias_exp, + down_bias=down_bias_exp, ) - # Unpad back to real token counts return _unpad(out_padded, counts, padded_counts, num_local_experts, permute_tokens.shape[0]) diff --git a/src/xorl/models/layers/moe/backend/quack.py b/src/xorl/models/layers/moe/backend/quack.py index d1daeefe..cf37b641 100644 --- a/src/xorl/models/layers/moe/backend/quack.py +++ b/src/xorl/models/layers/moe/backend/quack.py @@ -13,6 +13,8 @@ def quack_expert_forward( up_proj: torch.Tensor, down_proj: torch.Tensor, num_experts: int, + hidden_act: str = "silu", + gate_up_proj: torch.Tensor | None = None, **kwargs, ) -> torch.Tensor: """Forward pass using quack group GEMM kernels. @@ -27,11 +29,15 @@ def quack_expert_forward( up_proj: Up projection weights ``[num_experts, hidden, intermediate]``. down_proj: Down projection weights ``[num_experts, intermediate, hidden]``. num_experts: Total number of experts. - **kwargs: Extra arguments (ignored). + hidden_act: Activation kind ("silu" or "gelu_tanh"). + gate_up_proj: Optional pre-fused ``[num_experts, hidden, 2*intermediate]`` weight + (currently unused by the quack local path; accepted for interface parity). + **kwargs: Forwarded for forward compatibility; currently unused. Returns: Output tensor ``(num_tokens, hidden_dim)``. """ + del kwargs return quack_moe_forward( module=None, num_experts=num_experts, @@ -41,4 +47,6 @@ def quack_expert_forward( gate_proj=gate_proj, up_proj=up_proj, down_proj=down_proj, + gate_up_proj=gate_up_proj, + hidden_act=hidden_act, ) diff --git a/src/xorl/models/layers/moe/backend/triton.py b/src/xorl/models/layers/moe/backend/triton.py index ff9f971b..b32a465f 100644 --- a/src/xorl/models/layers/moe/backend/triton.py +++ b/src/xorl/models/layers/moe/backend/triton.py @@ -13,12 +13,14 @@ def triton_expert_forward( up_proj: torch.Tensor, down_proj: torch.Tensor, num_experts: int, + hidden_act: str = "silu", + gate_up_proj: torch.Tensor | None = None, **kwargs, ) -> torch.Tensor: """Forward pass using custom Triton group GEMM kernels. Uses ``xorl.ops.moe_experts_forward`` which dispatches to Triton kernels for - scatter/gather, group GEMM (``group_gemm_same_nk``), and fused SiLU+mul. + scatter/gather and group GEMM (``group_gemm_same_nk``). Args: hidden_states: Input tensor ``(num_tokens, hidden_dim)``. @@ -28,11 +30,15 @@ def triton_expert_forward( up_proj: Up projection weights ``[num_experts, hidden, intermediate]``. down_proj: Down projection weights ``[num_experts, intermediate, hidden]``. num_experts: Total number of experts. - **kwargs: Extra arguments (ignored). + hidden_act: Activation kind ("silu" or "gelu_tanh"). + gate_up_proj: Pre-fused ``[num_experts, hidden, 2*intermediate]`` weight + used by the fused-GEMM path (required by ``TritonMoeExpertsFunction``). + **kwargs: Forwarded for forward compatibility; currently unused. Returns: Output tensor ``(num_tokens, hidden_dim)``. """ + del kwargs return triton_moe_forward( module=None, num_experts=num_experts, @@ -42,4 +48,6 @@ def triton_expert_forward( gate_proj=gate_proj, up_proj=up_proj, down_proj=down_proj, + gate_up_proj=gate_up_proj, + hidden_act=hidden_act, ) diff --git a/src/xorl/models/layers/moe/experts.py b/src/xorl/models/layers/moe/experts.py index 86100c44..8e18caf8 100644 --- a/src/xorl/models/layers/moe/experts.py +++ b/src/xorl/models/layers/moe/experts.py @@ -39,6 +39,9 @@ class MoEExperts(nn.Module): ``gate_proj`` and ``up_proj`` are exposed as views into ``gate_up_proj`` for compatibility with existing backends and helpers. + Optional per-expert biases (``gate_up_bias``, ``down_bias``) default to + ``None`` and can be set by model-specific code (e.g. GPT-OSS). + Args: num_experts: Total number of experts. hidden_dim: Model hidden dimension. @@ -71,6 +74,15 @@ def __init__( requires_grad=True, ) self.act_fn = ACT2FN[hidden_act] + # String kind used by triton/native/quack backends (avoids name-sniffing). + from xorl.ops.moe.triton import normalize_hidden_act # noqa: PLC0415 + + self.hidden_act = normalize_hidden_act(hidden_act) + + # Optional per-expert biases (e.g. GPT-OSS). Set to actual tensors + # by model-specific code; None means no bias. + self.gate_up_bias = None + self.down_bias = None # EP dispatch strategy: "alltoall" (default) or "deepep" (NVLink-optimized) self.ep_dispatch: str = "alltoall" @@ -116,7 +128,9 @@ def forward( self.gate_proj.contiguous(), self.up_proj.contiguous(), self.down_proj, - self.act_fn, + hidden_act=self.hidden_act, + gate_up_bias=self.gate_up_bias, + down_bias=self.down_bias, ) # Check EP — use unified dispatch/compute/combine path @@ -140,6 +154,10 @@ def forward( up_proj, self.down_proj, num_experts=self.num_experts, + hidden_act=self.hidden_act, + gate_up_proj=self.gate_up_proj, + gate_up_bias=self.gate_up_bias, + down_bias=self.down_bias, ) @torch.compiler.disable @@ -235,6 +253,9 @@ def _ep_forward( self.down_proj, self.intermediate_size, expert_scores, + hidden_act=self.hidden_act, + gate_up_bias=self.gate_up_bias, + down_bias=self.down_bias, ) # Step 3: Combine expert outputs back to original ranks @@ -268,6 +289,9 @@ def _ep_forward_debug(self, dispatch_fn, combine_fn, compute_fn, dispatch_kwargs self.down_proj, self.intermediate_size, expert_scores, + hidden_act=self.hidden_act, + gate_up_bias=self.gate_up_bias, + down_bias=self.down_bias, ) ev[3].record() diff --git a/src/xorl/models/layers/moe/lora.py b/src/xorl/models/layers/moe/lora.py index 837ad78b..ba3ef6ea 100644 --- a/src/xorl/models/layers/moe/lora.py +++ b/src/xorl/models/layers/moe/lora.py @@ -81,6 +81,9 @@ def __init__( self.lora_config = lora_config or MoELoRAConfig() self.r = self.lora_config.r self.lora_alpha = self.lora_config.lora_alpha + self.active_r = self.r + self.active_lora_alpha = self.lora_alpha + self.use_rslora = self.lora_config.use_rslora # Base weights (frozen) in (G, K, N) format self.gate_up_proj = nn.Parameter( @@ -113,11 +116,7 @@ def __init__( self._create_lora_params("down_proj", num_exp, (1 if hybrid else num_exp), r, intermediate_size, hidden_dim) # Scaling factor - self.scaling = compute_lora_scaling( - self.lora_config.lora_alpha, - self.lora_config.r, - self.lora_config.use_rslora, - ) + self.scaling = compute_lora_scaling(self.lora_alpha, self.r, self.use_rslora) self.reset_lora_parameters() @@ -166,14 +165,28 @@ def reset_lora_parameters(self): nn.init.kaiming_uniform_(lora_A.data[i], a=math.sqrt(5)) nn.init.zeros_(lora_B.data) + def set_runtime_lora_config(self, lora_rank: int, lora_alpha: int) -> None: + """Update the active LoRA slice used during forward/merge/export.""" + if lora_rank <= 0 or lora_rank > self.r: + raise ValueError(f"Active LoRA rank must be in [1, {self.r}], got {lora_rank}") + self.active_r = lora_rank + self.active_lora_alpha = lora_alpha + + def _active_scaling(self) -> float: + return compute_lora_scaling(self.active_lora_alpha, self.active_r, self.use_rslora) + + def _active_lora_views(self, proj_name: str) -> tuple[torch.Tensor, torch.Tensor]: + lora_A = getattr(self, f"{proj_name}_lora_A")[..., : self.active_r].contiguous() + lora_B = getattr(self, f"{proj_name}_lora_B")[:, : self.active_r, ...].contiguous() + return lora_A, lora_B + def _compute_proj_delta(self, proj_name: str) -> torch.Tensor: """Compute LoRA delta for one projection. Returns [E, K, N] in GKN format.""" - lora_A = getattr(self, f"{proj_name}_lora_A") # [1 or E, in, r] - lora_B = getattr(self, f"{proj_name}_lora_B") # [E or 1, r, out] + lora_A, lora_B = self._active_lora_views(proj_name) E = max(lora_A.shape[0], lora_B.shape[0]) A = lora_A.expand(E, -1, -1) # [E, in, r] B = lora_B.expand(E, -1, -1) # [E, r, out] - return torch.bmm(A, B) * self.scaling # [E, in, out] = [E, K, N] + return torch.bmm(A, B) * self._active_scaling() # [E, in, out] = [E, K, N] def merge_weights(self) -> None: """Merge LoRA weights into base weights for inference. @@ -186,8 +199,9 @@ def merge_weights(self) -> None: if proj_name not in self.lora_config.target_modules: continue base = getattr(self, proj_name) - delta = self._compute_proj_delta(proj_name).to(base.dtype) - base.add_(delta) + delta = self._compute_proj_delta(proj_name) + merged = base.to(torch.float32) + delta + base.data.copy_(merged.to(base.dtype)) self.reset_lora_parameters() def forward( @@ -219,6 +233,9 @@ def forward( fn = MOE_EXPERT_BACKENDS_LORA[self.moe_implementation] gate_proj = self.gate_proj.contiguous() up_proj = self.up_proj.contiguous() + gate_proj_lora_A, gate_proj_lora_B = self._active_lora_views("gate_proj") + up_proj_lora_A, up_proj_lora_B = self._active_lora_views("up_proj") + down_proj_lora_A, down_proj_lora_B = self._active_lora_views("down_proj") return fn( num_experts=self.num_experts, routing_weights=routing_weights, @@ -227,13 +244,13 @@ def forward( 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, - up_proj_lora_A=self.up_proj_lora_A, - up_proj_lora_B=self.up_proj_lora_B, - down_proj_lora_A=self.down_proj_lora_A, - down_proj_lora_B=self.down_proj_lora_B, - scaling=self.scaling, + gate_proj_lora_A=gate_proj_lora_A, + gate_proj_lora_B=gate_proj_lora_B, + up_proj_lora_A=up_proj_lora_A, + up_proj_lora_B=up_proj_lora_B, + down_proj_lora_A=down_proj_lora_A, + down_proj_lora_B=down_proj_lora_B, + scaling=self._active_scaling(), ) def _ep_forward( @@ -264,6 +281,9 @@ def _ep_forward( compute_fn = EP_EXPERT_COMPUTE_LORA[self.moe_implementation] gate_proj = self.gate_proj.contiguous() up_proj = self.up_proj.contiguous() + gate_proj_lora_A, gate_proj_lora_B = self._active_lora_views("gate_proj") + up_proj_lora_A, up_proj_lora_B = self._active_lora_views("up_proj") + down_proj_lora_A, down_proj_lora_B = self._active_lora_views("down_proj") # Step 1: Dispatch tokens to expert-owning ranks dispatch_kwargs = self._build_dispatch_kwargs(hidden_states, routing_weights, selected_experts, parallel_state) @@ -276,18 +296,15 @@ def _ep_forward( 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, + gate_proj_lora_A, + gate_proj_lora_B, + up_proj_lora_A, + up_proj_lora_B, + down_proj_lora_A, + down_proj_lora_B, + self._active_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) @@ -339,32 +356,36 @@ def _eager_lora_forward(self, hidden_states: torch.Tensor, expert_idx: int) -> t # x @ W — no transpose needed with (G, K, N) format gate_proj_out = torch.matmul(hidden_states, self.gate_proj[expert_idx]) up_proj_out = torch.matmul(hidden_states, self.up_proj[expert_idx]) + active_scaling = self._active_scaling() if "gate_proj" in self.lora_config.target_modules: - A = self.gate_proj_lora_A[min(expert_idx, self.gate_proj_lora_A.shape[0] - 1)].to(compute_dtype) - B = self.gate_proj_lora_B[expert_idx].to(compute_dtype) - gate_proj_out = gate_proj_out + torch.matmul(torch.matmul(hidden_states, A), B) * self.scaling + gate_A, gate_B = self._active_lora_views("gate_proj") + A = gate_A[min(expert_idx, gate_A.shape[0] - 1)].to(compute_dtype) + B = gate_B[expert_idx].to(compute_dtype) + gate_proj_out = gate_proj_out + torch.matmul(torch.matmul(hidden_states, A), B) * active_scaling if "up_proj" in self.lora_config.target_modules: - A = self.up_proj_lora_A[min(expert_idx, self.up_proj_lora_A.shape[0] - 1)].to(compute_dtype) - B = self.up_proj_lora_B[expert_idx].to(compute_dtype) - up_proj_out = up_proj_out + torch.matmul(torch.matmul(hidden_states, A), B) * self.scaling + up_A, up_B = self._active_lora_views("up_proj") + A = up_A[min(expert_idx, up_A.shape[0] - 1)].to(compute_dtype) + B = up_B[expert_idx].to(compute_dtype) + up_proj_out = up_proj_out + torch.matmul(torch.matmul(hidden_states, A), B) * active_scaling out = self.act_fn(gate_proj_out) * up_proj_out down_out = torch.matmul(out, self.down_proj[expert_idx]) if "down_proj" in self.lora_config.target_modules: - A = self.down_proj_lora_A[expert_idx].to(compute_dtype) - B = self.down_proj_lora_B[min(expert_idx, self.down_proj_lora_B.shape[0] - 1)].to(compute_dtype) - down_out = down_out + torch.matmul(torch.matmul(out, A), B) * self.scaling + down_A, down_B = self._active_lora_views("down_proj") + A = down_A[expert_idx].to(compute_dtype) + B = down_B[min(expert_idx, down_B.shape[0] - 1)].to(compute_dtype) + down_out = down_out + torch.matmul(torch.matmul(out, A), B) * active_scaling return down_out def extra_repr(self) -> str: return ( f"num_experts={self.num_experts}, hidden_dim={self.hidden_dim}, " - f"intermediate_size={self.intermediate_size}, r={self.lora_config.r}, " - f"lora_alpha={self.lora_config.lora_alpha}, " + f"intermediate_size={self.intermediate_size}, r={self.active_r}, max_r={self.r}, " + f"lora_alpha={self.active_lora_alpha}, " f"target_modules={self.lora_config.target_modules}" ) diff --git a/src/xorl/models/layers/moe/moe_block.py b/src/xorl/models/layers/moe/moe_block.py index c33acc90..c99286c5 100644 --- a/src/xorl/models/layers/moe/moe_block.py +++ b/src/xorl/models/layers/moe/moe_block.py @@ -7,7 +7,7 @@ import torch.nn.functional as F from .experts import MoEExperts -from .router import TopKRouter +from .router import TopKRouter, balanced_synthetic_routing from .routing_replay import RoutingReplay, get_replay_stage @@ -50,6 +50,7 @@ def __init__( norm_topk_prob: bool = True, moe_implementation: str = "triton", train_router: bool = True, + record_routing_weights: bool = True, ): super().__init__() self.num_experts = num_experts @@ -58,6 +59,7 @@ def __init__( self.intermediate_size = intermediate_size self.moe_implementation = moe_implementation self.train_router = train_router + self.record_routing_weights = record_routing_weights # Gate linear — directly on this module for checkpoint path ``mlp.gate.weight`` self.gate = nn.Linear(hidden_size, num_experts, bias=False) @@ -121,6 +123,16 @@ def _regather_routing(self, router_logits, cached_experts, input_dtype): graph for gate gradients), then gather with cached indices. Gradients flow through softmax -> gate naturally, no straight-through hack needed. """ + if self.router.synthetic_routing_mode == "balanced": + routing_weights, _ = balanced_synthetic_routing( + cached_experts.size(0), + self.num_experts, + cached_experts.size(1), + router_logits.device, + input_dtype, + ) + return cached_experts, routing_weights + softmax_probs = F.softmax(router_logits, dim=1, dtype=torch.float) routing_weights = torch.gather(softmax_probs, 1, cached_experts) if self.router.norm_topk_prob: @@ -182,23 +194,24 @@ def route(self, hidden_states: torch.Tensor): 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) + if self.record_routing_weights: + 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) @@ -292,4 +305,5 @@ def from_config(cls, config, moe_implementation: str = "triton"): norm_topk_prob=config.norm_topk_prob, moe_implementation=moe_implementation, train_router=getattr(config, "train_router", False), + record_routing_weights=getattr(config, "record_routing_weights", True), ) diff --git a/src/xorl/models/layers/moe/router.py b/src/xorl/models/layers/moe/router.py index 6b169fe2..2f88d196 100644 --- a/src/xorl/models/layers/moe/router.py +++ b/src/xorl/models/layers/moe/router.py @@ -1,10 +1,41 @@ """Token-choice top-k router for MoE layers.""" +import os + import torch import torch.nn as nn import torch.nn.functional as F +_SYNTHETIC_ROUTING_ENV = "XORL_MOE_SYNTHETIC_ROUTING" + + +def _get_synthetic_routing_mode() -> str: + mode = os.environ.get(_SYNTHETIC_ROUTING_ENV, "").strip().lower() + if mode in {"", "none", "default"}: + return "" + if mode == "balanced": + return mode + raise ValueError(f"Unsupported {_SYNTHETIC_ROUTING_ENV}={mode!r}; expected 'balanced'.") + + +def balanced_synthetic_routing( + num_tokens: int, + num_experts: int, + top_k: int, + device: torch.device, + dtype: torch.dtype, +) -> tuple[torch.Tensor, torch.Tensor]: + if top_k > num_experts: + raise ValueError(f"top_k ({top_k}) must be <= num_experts ({num_experts})") + + token_offsets = torch.arange(num_tokens, device=device, dtype=torch.long).unsqueeze(1) * top_k + topk_offsets = torch.arange(top_k, device=device, dtype=torch.long).unsqueeze(0) + selected_experts = (token_offsets + topk_offsets) % num_experts + routing_weights = torch.full((num_tokens, top_k), 1.0 / top_k, device=device, dtype=dtype) + return routing_weights, selected_experts + + class TopKRouter(nn.Module): """Top-K routing: softmax -> topk -> optional renormalization. @@ -28,6 +59,7 @@ def __init__( self.num_experts = num_experts self.top_k = top_k self.norm_topk_prob = norm_topk_prob + self.synthetic_routing_mode = _get_synthetic_routing_mode() def forward( self, @@ -44,6 +76,15 @@ def forward( routing_weights: ``(num_tokens, top_k)`` weights per selected expert. selected_experts: ``(num_tokens, top_k)`` expert indices. """ + if self.synthetic_routing_mode == "balanced": + return balanced_synthetic_routing( + router_logits.size(0), + self.num_experts, + self.top_k, + router_logits.device, + input_dtype, + ) + routing_weights = F.softmax(router_logits, dim=1, dtype=torch.float) routing_weights, selected_experts = torch.topk(routing_weights, self.top_k, dim=-1) if self.norm_topk_prob: diff --git a/src/xorl/models/layers/moe/routing_replay.py b/src/xorl/models/layers/moe/routing_replay.py index 9c7ae3bb..d664ffc4 100644 --- a/src/xorl/models/layers/moe/routing_replay.py +++ b/src/xorl/models/layers/moe/routing_replay.py @@ -40,34 +40,39 @@ class RoutingReplay: def __init__(self): self.forward_index: int = 0 self.backward_index: int = 0 - self.top_indices_list: List[torch.Tensor] = [] # CPU pinned - self.top_weights_list: List[torch.Tensor] = [] # CPU pinned routing weights (R3 logits) + self.top_indices_list: List[torch.Tensor] = [] # CPU buffers + self.top_weights_list: List[torch.Tensor] = [] # CPU routing weight buffers (R3 logits) RoutingReplay._instances.append(self) @torch.compiler.disable def record(self, selected_experts: torch.Tensor): - """Append routing decision (CPU pinned copy). + """Append routing decision (CPU copy). Disabled for torch.compile — pin_memory and list-append side effects are not supported by Inductor/Dynamo. """ - buf = torch.empty_like(selected_experts, device="cpu", pin_memory=True) + buf = torch.empty_like(selected_experts, device="cpu", pin_memory=torch.cuda.is_available()) buf.copy_(selected_experts) self.top_indices_list.append(buf) @torch.compiler.disable def record_weights(self, routing_weights: torch.Tensor): - """Append routing weights (CPU pinned copy) for R3 weight replay.""" - buf = torch.empty_like(routing_weights, device="cpu", pin_memory=True) + """Append routing weights (CPU copy) for R3 weight replay.""" + buf = torch.empty_like(routing_weights, device="cpu", pin_memory=torch.cuda.is_available()) buf.copy_(routing_weights) self.top_weights_list.append(buf) + def _target_device(self) -> torch.device: + if torch.cuda.is_available(): + return torch.device("cuda", torch.cuda.current_device()) + return torch.device("cpu") + @torch.compiler.disable def pop_forward(self) -> torch.Tensor: """Read routing for forward replay, advance forward_index.""" idx = self.top_indices_list[self.forward_index] self.forward_index += 1 - return idx.to(torch.cuda.current_device(), non_blocking=True) + return idx.to(self._target_device(), non_blocking=torch.cuda.is_available()) @torch.compiler.disable def pop_forward_weights(self) -> Optional[torch.Tensor]: @@ -75,21 +80,25 @@ 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( + self._target_device(), non_blocking=torch.cuda.is_available() + ) @torch.compiler.disable def pop_backward(self) -> torch.Tensor: """Read routing for checkpoint recompute, advance backward_index.""" idx = self.top_indices_list[self.backward_index] self.backward_index += 1 - return idx.to(torch.cuda.current_device(), non_blocking=True) + return idx.to(self._target_device(), non_blocking=torch.cuda.is_available()) @torch.compiler.disable 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( + self._target_device(), non_blocking=torch.cuda.is_available() + ) @property def has_weights(self) -> bool: diff --git a/src/xorl/models/layers/normalization.py b/src/xorl/models/layers/normalization.py index 7a0552cd..69de3223 100644 --- a/src/xorl/models/layers/normalization.py +++ b/src/xorl/models/layers/normalization.py @@ -8,6 +8,7 @@ RMSNormMode = Literal["eager", "native", "compile"] _RMSNORM_MODE: RMSNormMode = "native" _COMPILED_NATIVE_RMS_NORM: Optional[Callable[[torch.Tensor, torch.Tensor, float], torch.Tensor]] = None +_COMPILED_EAGER_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 @@ -42,6 +43,15 @@ def compiled_rms_norm(hidden_states: torch.Tensor, weight: torch.Tensor, varianc return _COMPILED_NATIVE_RMS_NORM(hidden_states, weight, variance_epsilon) +def compiled_eager_rms_norm(hidden_states: torch.Tensor, weight: torch.Tensor, variance_epsilon: float) -> torch.Tensor: + """Compiled fp32-upcast eager RMSNorm. Use when ``F.rms_norm``'s bf16 path is + too imprecise (e.g., GLM-4 MoE 92+ norm layers compounding bf16 error).""" + global _COMPILED_EAGER_RMS_NORM + if _COMPILED_EAGER_RMS_NORM is None: + _COMPILED_EAGER_RMS_NORM = torch.compile(eager_rms_norm) + return _COMPILED_EAGER_RMS_NORM(hidden_states, weight, variance_epsilon) + + def eager_zero_centered_rms_norm( hidden_states: torch.Tensor, weight: torch.Tensor, diff --git a/src/xorl/models/layers/rope.py b/src/xorl/models/layers/rope.py index 94eafe71..e7bb319a 100644 --- a/src/xorl/models/layers/rope.py +++ b/src/xorl/models/layers/rope.py @@ -477,11 +477,19 @@ def _naive_apply_rotary_pos_emb(q, k, cos, sin): """Naive RoPE application (pure PyTorch, no fused kernel). All tensors use [B, S, H, D] layout. cos/sin are [B, S, D]. + Handles partial rotary automatically when cos/sin dim < head_dim. """ cos = cos.unsqueeze(2) sin = sin.unsqueeze(2) - q_embed = (q * cos) + (rotate_half(q) * sin) - k_embed = (k * cos) + (rotate_half(k) * sin) + rotary_dim = cos.shape[-1] + if q.shape[-1] > rotary_dim: + q_rot, q_pass = q[..., :rotary_dim], q[..., rotary_dim:] + k_rot, k_pass = k[..., :rotary_dim], k[..., rotary_dim:] + q_embed = torch.cat([(q_rot * cos) + (rotate_half(q_rot) * sin), q_pass], dim=-1) + k_embed = torch.cat([(k_rot * cos) + (rotate_half(k_rot) * sin), k_pass], dim=-1) + else: + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed diff --git a/src/xorl/models/loader.py b/src/xorl/models/loader.py index 594d3723..d2fca3b2 100644 --- a/src/xorl/models/loader.py +++ b/src/xorl/models/loader.py @@ -2,10 +2,6 @@ import torch import torch.nn as nn -from transformers import ( - AutoModel, - AutoModelForCausalLM, -) from ..utils import logging from .module_utils import all_ranks_load_weights, init_empty_weights @@ -16,11 +12,10 @@ class ModelLoader: - """Unified model loader for both HuggingFace and custom xorl models. + """Model loader for xorl-registered architectures. - Takes a model factory callable (e.g., ``AutoModelForCausalLM.from_config`` - or ``model_cls._from_config``) and handles meta-init, device placement, - and weight loading. + Takes a model factory callable (e.g., ``model_cls._from_config``) and + handles meta-init, device placement, and weight loading. """ def __init__(self, model_factory: Callable[..., nn.Module], description: str = ""): @@ -69,7 +64,15 @@ def get_loader(model_config) -> ModelLoader: model_cls = ModelRegistry.get_model_cls_from_model_arch(model_arch) return ModelLoader(model_cls._from_config, description=f"xorl/{model_arch}") - if "ForCausalLM" in model_arch and type(model_config) in AutoModelForCausalLM._model_mapping.keys(): - return ModelLoader(AutoModelForCausalLM.from_config, description=f"huggingface/{model_arch}") + if ModelRegistry.import_errors: + failures = "\n".join(f" - {name}: {err!r}" for name, err in ModelRegistry.import_errors.items()) + raise NotImplementedError( + f"Architecture {model_arch!r} is not registered in xorl, and the following " + f"xorl modeling modules failed to import (one of them may provide this arch):\n" + f"{failures}" + ) - return ModelLoader(AutoModel.from_config, description=f"huggingface/{model_arch}") + raise NotImplementedError( + f"Architecture {model_arch!r} is not implemented in xorl. Supported architectures: " + f"{sorted(ModelRegistry.supported_models)}" + ) diff --git a/src/xorl/models/module_utils.py b/src/xorl/models/module_utils.py index 6ac45c14..bee30f69 100644 --- a/src/xorl/models/module_utils.py +++ b/src/xorl/models/module_utils.py @@ -1,6 +1,7 @@ import json import math import os +import pickle import re import time from collections import OrderedDict, deque @@ -16,16 +17,22 @@ from safetensors.torch import load_file, save_file from torch import distributed as dist from torch import nn +from torch.distributed import ProcessGroup from torch.distributed._tensor import Shard as DTShard +from torch.distributed.device_mesh import DeviceMesh +from torch.distributed.tensor import DTensor, Replicate from tqdm import tqdm +from transformers import PretrainedConfig, PreTrainedTokenizerBase 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 xorl.distributed.parallel_state import get_parallel_state from xorl.lora.modules.linear import LoraLinear from xorl.models.checkpoint_handlers.buffers import ( # noqa: F401 + FUSED_EXPERT_PATTERN, ExpertWeightBuffer, checkpoint_has_per_expert_weights, + parse_expert_full_key, parse_expert_key, ) from xorl.ops.loss import get_loss_function @@ -45,21 +52,15 @@ logger = logging.get_logger(__name__) -_cpu_group = None _weight_load_group = None - - -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 +_grouped_weight_load_group = None +_grouped_weight_load_group_ranks: Optional[Tuple[int, ...]] = None +_grouped_dense_weight_load_group = None +_grouped_dense_weight_load_group_ranks: Optional[Tuple[int, ...]] = None +_save_sync_group = None +_save_sync_group_backend: Optional[str] = None +_cpu_save_device_mesh_cache: Dict[Tuple[str, Tuple[int, ...], Tuple[int, ...], Tuple[str, ...]], DeviceMesh] = {} +_UNSUPPORTED_DTENSOR_ROOT_GATHER = object() def _get_weight_load_group(): @@ -76,24 +77,406 @@ def _get_weight_load_group(): 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: +def _get_object_broadcast_device(group) -> Optional[torch.device]: + """Return the device to use for object broadcasts on the given process group.""" + if not dist.is_available() or not dist.is_initialized(): return None - if dist.get_backend(group) == "nccl": + + backend = dist.get_backend() if group is None else dist.get_backend(group) + if backend == "nccl": return torch.device(f"{get_device_type()}:{get_device_id()}") return None +def _get_cpu_save_device_mesh(device_mesh: DeviceMesh) -> DeviceMesh: + """Mirror a DTensor device mesh onto CPU/Gloo for low-memory save gathers.""" + mesh_tensor = device_mesh.mesh.detach().cpu().contiguous() + mesh_dim_names = tuple(device_mesh.mesh_dim_names or tuple(f"dim_{idx}" for idx in range(mesh_tensor.ndim))) + cache_key = ( + device_mesh.device_type, + tuple(mesh_tensor.size()), + tuple(int(rank) for rank in mesh_tensor.reshape(-1).tolist()), + mesh_dim_names, + ) + cached_mesh = _cpu_save_device_mesh_cache.get(cache_key) + if cached_mesh is None: + backend_override = tuple(("gloo", None) for _ in range(mesh_tensor.ndim)) + cached_mesh = DeviceMesh( + device_type="cpu", + mesh=mesh_tensor, + mesh_dim_names=mesh_dim_names, + backend_override=backend_override, + ) + _cpu_save_device_mesh_cache[cache_key] = cached_mesh + return cached_mesh + + +def _get_mesh_group_ranks(device_mesh: DeviceMesh, mesh_dim: int) -> Tuple[int, ...]: + """Return global ranks in the current rank's mesh-dim subgroup.""" + mesh_tensor = device_mesh.mesh.detach().cpu() + coord = device_mesh.get_coordinate() + if coord is None: + raise RuntimeError("Current rank is not part of the provided device mesh.") + + index = [] + for dim_idx, coord_idx in enumerate(coord): + index.append(slice(None) if dim_idx == mesh_dim else coord_idx) + group_ranks = mesh_tensor[tuple(index)].reshape(-1).tolist() + return tuple(int(rank) for rank in group_ranks) + + +def _get_rank_coordinate(device_mesh: DeviceMesh, global_rank: int) -> Optional[Tuple[int, ...]]: + """Return the coordinate of a global rank in a device mesh.""" + mesh_tensor = device_mesh.mesh.detach().cpu() + matches = (mesh_tensor == global_rank).nonzero(as_tuple=False) + if matches.numel() == 0: + return None + return tuple(int(coord) for coord in matches[0].tolist()) + + +def _get_mesh_group_root_rank(device_mesh: DeviceMesh, mesh_dim: int, target_index: int) -> int: + """Return the global rank in the current mesh-dim subgroup at ``target_index``.""" + mesh_tensor = device_mesh.mesh.detach().cpu() + coord = device_mesh.get_coordinate() + if coord is None: + raise RuntimeError("Current rank is not part of the provided device mesh.") + index = list(coord) + index[mesh_dim] = target_index + return int(mesh_tensor[tuple(index)].item()) + + +def _gather_sharded_tensor_to_group_root( + local_tensor: "torch.Tensor", + *, + tensor_dim: int, + full_dim_size: int, + group, + group_ranks: Sequence[int], + dst_rank: int, +) -> Optional["torch.Tensor"]: + """Gather shards from one mesh-dim subgroup to a single root rank.""" + if not group_ranks: + return None + + local_cpu = local_tensor.detach().cpu().contiguous() + if dst_rank not in group_ranks: + group_dst = 0 + keep_result = False + else: + group_dst = group_ranks.index(dst_rank) + keep_result = dist.get_rank() == dst_rank + + local_group_rank = group_ranks.index(dist.get_rank()) + chunk_size = math.ceil(full_dim_size / len(group_ranks)) + padded_shape = list(local_cpu.shape) + padded_shape[tensor_dim] = chunk_size + padded = local_cpu.new_zeros(padded_shape) + if local_cpu.shape[tensor_dim] > 0: + index = [slice(None)] * local_cpu.ndim + index[tensor_dim] = slice(0, local_cpu.shape[tensor_dim]) + padded[tuple(index)] = local_cpu + + gather_list = None + if local_group_rank == group_dst: + gather_list = [padded.new_empty(padded_shape) for _ in range(len(group_ranks))] + + dist.gather(padded, gather_list=gather_list, group=group, group_dst=group_dst) + + if local_group_rank != group_dst: + return None + + shards = [] + for shard_idx, gathered in enumerate(gather_list or []): + start = shard_idx * chunk_size + remaining = max(full_dim_size - start, 0) + take = min(chunk_size, remaining) + if take <= 0: + continue + if take != chunk_size: + gathered = gathered.narrow(tensor_dim, 0, take) + shards.append(gathered) + + if not shards: + result = padded.new_empty(padded_shape) + result = result.narrow(tensor_dim, 0, 0) + elif len(shards) == 1: + result = shards[0] + else: + result = torch.cat(shards, dim=tensor_dim) + + return result if keep_result else None + + +def _gather_dtensor_to_rank(raw_tensor: DTensor, dst_rank: int) -> Union["torch.Tensor", None, object]: + """Gather a DTensor to one rank on CPU/Gloo when placements allow it.""" + placements = tuple(raw_tensor.placements) + shard_mesh_dims = [mesh_dim for mesh_dim, placement in enumerate(placements) if isinstance(placement, DTShard)] + if len(shard_mesh_dims) == 0: + local_tensor = raw_tensor.to_local() + if hasattr(local_tensor, "wait"): + local_tensor = local_tensor.wait() + return local_tensor.detach().cpu() if dist.get_rank() == dst_rank else None + + cpu_mesh = _get_cpu_save_device_mesh(raw_tensor.device_mesh) + dst_coord = _get_rank_coordinate(cpu_mesh, dst_rank) + if dst_coord is None: + return _UNSUPPORTED_DTENSOR_ROOT_GATHER + + local_tensor = raw_tensor.to_local() + if hasattr(local_tensor, "wait"): + local_tensor = local_tensor.wait() + partial = local_tensor.detach().cpu().contiguous() + full_shape = tuple(raw_tensor.size()) + + for shard_mesh_dim in reversed(shard_mesh_dims): + placement = placements[shard_mesh_dim] + if not isinstance(placement, DTShard): + return _UNSUPPORTED_DTENSOR_ROOT_GATHER + group = cpu_mesh.get_group(shard_mesh_dim if cpu_mesh.ndim > 1 else None) + group_ranks = _get_mesh_group_ranks(cpu_mesh, shard_mesh_dim if cpu_mesh.ndim > 1 else 0) + if not group_ranks: + return _UNSUPPORTED_DTENSOR_ROOT_GATHER + + if len(shard_mesh_dims) == 1: + stage_dst_rank = dst_rank + else: + stage_dst_rank = _get_mesh_group_root_rank(cpu_mesh, shard_mesh_dim, dst_coord[shard_mesh_dim]) + + partial = _gather_sharded_tensor_to_group_root( + partial, + tensor_dim=placement.dim, + full_dim_size=full_shape[placement.dim], + group=group, + group_ranks=group_ranks, + dst_rank=stage_dst_rank, + ) + + coord = cpu_mesh.get_coordinate() + if coord is None: + return _UNSUPPORTED_DTENSOR_ROOT_GATHER + if coord[shard_mesh_dim] != dst_coord[shard_mesh_dim]: + return None + if partial is None: + return None + + return partial if dist.get_rank() == dst_rank else None + + +def _materialize_tensor_for_save( + tensor: "torch.Tensor", + dst_rank: Optional[int] = None, +) -> Optional["torch.Tensor"]: + """Materialize a parameter/buffer for save without gathering full DTensors on CUDA.""" + raw_tensor = tensor.data if hasattr(tensor, "data") else tensor + if isinstance(raw_tensor, DTensor): + if dst_rank is not None: + gathered = _gather_dtensor_to_rank(raw_tensor, dst_rank) + if gathered is not _UNSUPPORTED_DTENSOR_ROOT_GATHER: + return gathered + logger.warning_once( + "Falling back to replicated DTensor save materialization for unsupported placements: " + f"{raw_tensor.placements}" + ) + + local_tensor = raw_tensor.to_local() + if hasattr(local_tensor, "wait"): + local_tensor = local_tensor.wait() + local_cpu = local_tensor.detach().cpu() + cpu_mesh = _get_cpu_save_device_mesh(raw_tensor.device_mesh) + cpu_dtensor = DTensor.from_local( + local_cpu, + device_mesh=cpu_mesh, + placements=raw_tensor.placements, + run_check=False, + shape=raw_tensor.size(), + stride=raw_tensor.stride(), + ) + replicated = cpu_dtensor.redistribute( + device_mesh=cpu_mesh, + placements=[Replicate() for _ in cpu_dtensor.placements], + ).to_local() + if hasattr(replicated, "wait"): + replicated = replicated.wait() + return replicated + + if dst_rank is not None and dist.is_available() and dist.is_initialized() and dist.get_rank() != dst_rank: + return None + + if hasattr(raw_tensor, "full_tensor"): + return raw_tensor.full_tensor() + + return raw_tensor + + 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) + _broadcast_object_list(obj_list, src=src, group=group) + + +def _broadcast_object_list(obj_list: List[Any], src: int, group) -> None: + """Broadcast Python metadata on the provided process group.""" + device = _get_object_broadcast_device(group) + if device is None: + dist.broadcast_object_list(obj_list, src=src, group=group) + return + + is_source = dist.get_rank() == src + payload = pickle.dumps(obj_list, protocol=pickle.HIGHEST_PROTOCOL) if is_source else b"" + size_tensor = torch.tensor([len(payload)], dtype=torch.int64, device=device) + dist.broadcast(size_tensor, src=src, group=group) + + if is_source: + payload_tensor = torch.tensor(list(payload), dtype=torch.uint8, device=device) + else: + payload_tensor = torch.empty(int(size_tensor.item()), dtype=torch.uint8, device=device) + + dist.broadcast(payload_tensor, src=src, group=group) + + if not is_source: + obj_list[:] = pickle.loads(bytes(payload_tensor.cpu().tolist())) + + +def _normalize_checkpoint_key_for_filter(key: str) -> Optional[str]: + """Normalize raw checkpoint keys for lightweight load-time filtering.""" + if key.startswith("vision_tower.") or key.startswith("mm_projector."): + return None + if key.startswith("model.language_model."): + return "model." + key.removeprefix("model.language_model.") + if key.startswith("language_model."): + return key.removeprefix("language_model.") + return key + + +def _matches_checkpoint_skip_key_pattern(key: str, model: object) -> bool: + """Return True when a raw or converted checkpoint key is model-declared skip state.""" + for pattern in getattr(model, "_checkpoint_skip_key_patterns", ()): + if re.match(pattern, key): + return True + return False + + +_FUSED_EXPERT_CHECKPOINT_PATTERN = re.compile(r"^model\.layers\.\d+\.mlp\.experts\.(gate_up|down)_proj(?:\..+)?$") +_FFN_EXPERT_CHECKPOINT_PATTERN = re.compile(r"^(?:model\.)?layers\.\d+\.ffn\.experts\.\d+\.w[123]\.(?:weight|scale)$") + + +def _is_checkpoint_expert_key(key: str) -> bool: + """Return True when a raw checkpoint key belongs to MoE expert weights.""" + normalized = _normalize_checkpoint_key_for_filter(key) + if normalized is None: + return False + return ( + parse_expert_full_key(normalized) is not None + or FUSED_EXPERT_PATTERN.match(normalized) is not None + or _FUSED_EXPERT_CHECKPOINT_PATTERN.match(normalized) is not None + or _FFN_EXPERT_CHECKPOINT_PATTERN.match(normalized) is not None + ) + + +def _is_expert_parameter_name(parameter_name: str, parallel_plan: Optional["ParallelPlan"]) -> bool: + """Return True when a model parameter is EP-sharded expert state.""" + if parallel_plan is not None: + is_expert = getattr(parallel_plan, "is_expert_parameter", None) + if callable(is_expert): + return bool(is_expert(parameter_name)) + private_is_expert = getattr(parallel_plan, "_is_expert_parameter", None) + if callable(private_is_expert): + return bool(private_is_expert(parameter_name)) + return FUSED_EXPERT_PATTERN.match(parameter_name) is not None + + +def _get_grouped_weight_load_group(parallel_state) -> Optional[ProcessGroup]: + """Return the subgroup used by grouped weight loading. + + Grouped loading is only meaningful when EP is enabled: one leader per + ``ep_fsdp`` group reads and fan-outs tensors to the ranks that share that + expert shard. This keeps expert weights correct without forcing a global + rank-0 to materialize every expert tensor. Use a dedicated process group + with the weight-load timeout so long checkpoint reads do not inherit the + shorter default timeout from the DeviceMesh subgroup. + """ + global _grouped_weight_load_group, _grouped_weight_load_group_ranks + if _grouped_weight_load_group is not None: + return _grouped_weight_load_group + if ( + parallel_state is None + or not getattr(parallel_state, "ep_enabled", False) + or getattr(parallel_state, "ep_fsdp_device_mesh", None) is None + ): + return None + + if not dist.is_available() or not dist.is_initialized(): + return None + + mesh_tensor = parallel_state.ep_fsdp_device_mesh.mesh.detach().cpu().contiguous() + if mesh_tensor.ndim == 0: + return None + + timeout_sec = int(os.getenv("XORL_WEIGHT_LOAD_TIMEOUT_SEC", "7200")) + backend = dist.get_backend() + rank = dist.get_rank() + created_group = None + created_ranks = None + for group_ranks_tensor in mesh_tensor.view(-1, mesh_tensor.size(-1)): + ranks = tuple(int(member_rank) for member_rank in group_ranks_tensor.tolist()) + group = dist.new_group(ranks=list(ranks), backend=backend, timeout=timedelta(seconds=timeout_sec)) + if rank in ranks: + created_group = group + created_ranks = ranks + + _grouped_weight_load_group = created_group + _grouped_weight_load_group_ranks = created_ranks + return _grouped_weight_load_group + + +def _get_grouped_weight_load_prefetch_count(shard_count: int) -> int: + """Prefetch depth for grouped loading. + + We intentionally keep this small. Grouped loading already reduces reader + fan-out, so large per-rank prefetch would just recreate the WekaFS stampede + that motivated this mode. + """ + configured = int(os.getenv("XORL_GROUPED_WEIGHT_LOAD_PREFETCH_COUNT", "2")) + return max(1, min(configured, shard_count)) + + +def _get_grouped_dense_weight_load_group(): + """Return the per-node process group used for grouped dense/shared replication. + + Grouped loading already uses EP-FSDP subgroups for expert tensors. Dense and + shared weights should not use an all-world metadata broadcast, but loading + them on every rank still creates avoidable checkpoint I/O. We instead create + one node-local group per host and broadcast dense/shared tensors inside that + group so each shard is read once per node rather than once per rank. + """ + global _grouped_dense_weight_load_group, _grouped_dense_weight_load_group_ranks + if _grouped_dense_weight_load_group is not None: + return _grouped_dense_weight_load_group + + if not dist.is_available() or not dist.is_initialized(): + return None + + world_size = dist.get_world_size() + local_world_size = max(1, int(os.environ.get("LOCAL_WORLD_SIZE", "1"))) + if world_size <= 1 or local_world_size <= 1: + _grouped_dense_weight_load_group_ranks = (dist.get_rank(),) + return None + + timeout_sec = int(os.getenv("XORL_WEIGHT_LOAD_TIMEOUT_SEC", "7200")) + rank = dist.get_rank() + created_group = None + created_ranks = None + backend = dist.get_backend() + for start_rank in range(0, world_size, local_world_size): + ranks = tuple(range(start_rank, min(start_rank + local_world_size, world_size))) + group = dist.new_group(ranks=list(ranks), backend=backend, timeout=timedelta(seconds=timeout_sec)) + if rank in ranks: + created_group = group + created_ranks = ranks + + _grouped_dense_weight_load_group = created_group + _grouped_dense_weight_load_group_ranks = created_ranks + return _grouped_dense_weight_load_group def _get_checkpoint_keys(weights_path: str) -> Optional[Set[str]]: @@ -358,6 +741,11 @@ def _try_load_state_dict(weights_path: str, **kwargs): # Single rank: do everything locally (no broadcast needed) return _try_load_state_dict_local(weights_path, **kwargs) + if os.path.isdir(weights_path): + # Shared local snapshots are cheap to resolve independently and avoid + # another cross-rank metadata broadcast before tensor loading starts. + return _try_load_state_dict_local(weights_path, **kwargs) + # Multi-rank: rank 0 resolves paths, broadcasts to all resolved_paths = [None] if rank == 0: @@ -505,6 +893,26 @@ def _build_expert_scatter_list( return full_tensor, scatter_list +def _build_group_scatter_list( + tensor: "torch.Tensor", + target_shape: Tuple[int, ...], + group_size: int, + torch_device: "torch.device", +) -> Tuple["torch.Tensor", List["torch.Tensor"]]: + """Prepare source-side per-rank views for subgroup 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) + + local_rows = target_shape[0] + scatter_list: List[torch.Tensor] = [None] * group_size + for group_rank in range(group_size): + scatter_list[group_rank] = full_tensor.narrow(0, group_rank * local_rows, local_rows) + + return full_tensor, scatter_list + + class _MultiStreamDMA: """Manages multi-stream H2D DMA transfers for overlapping copy engine usage. @@ -578,6 +986,66 @@ def flush(self) -> None: _active_dma_scheduler: Optional[_MultiStreamDMA] = None +def _copy_into_existing_dtensor_shard(dtensor: "torch.Tensor", tensor: "torch.Tensor") -> bool: + """Copy a full tensor directly into a 1D DTensor's local shard. + + When the caller already holds the full parameter value on every rank + (broadcast/grouped load modes), rebuilding a DTensor via ``distribute_tensor`` + adds extra padding logic and collectives that are not needed. This helper + handles the common 1D-mesh cases directly: + + - replicated DTensors: copy the full tensor into ``_local_tensor`` + - singly sharded DTensors: split the full tensor locally and copy the + current rank's shard into ``_local_tensor`` + """ + if not hasattr(dtensor, "_local_tensor"): + return False + + device_mesh = getattr(dtensor, "device_mesh", None) + placements = getattr(dtensor, "placements", None) + if device_mesh is None or placements is None: + return False + if getattr(device_mesh, "ndim", None) != 1: + return False + + shard_placements = [placement for placement in placements if isinstance(placement, DTShard)] + if len(shard_placements) > 1: + return False + + local_tensor = dtensor._local_tensor + if len(shard_placements) == 0: + if tuple(local_tensor.shape) != tuple(tensor.shape): + return False + local_tensor.copy_(tensor.to(device=local_tensor.device, dtype=local_tensor.dtype)) + return True + + shard = shard_placements[0] + mesh_size = device_mesh.size() + local_rank = device_mesh.get_local_rank() + shards, _ = shard._split_tensor(tensor, mesh_size, with_padding=True, contiguous=True) + local_shard = shards[local_rank] + expected_shape = tuple(local_tensor.shape) + if tuple(local_shard.shape) != expected_shape: + if local_shard.ndim != local_tensor.ndim: + return False + can_trim = True + for dim, (actual, expected) in enumerate(zip(local_shard.shape, expected_shape)): + if dim == shard.dim: + if expected > actual: + can_trim = False + break + elif expected != actual: + can_trim = False + break + if not can_trim: + return False + local_shard = local_shard.narrow(shard.dim, 0, expected_shape[shard.dim]) + + local_shard = local_shard.to(device=local_tensor.device, dtype=local_tensor.dtype) + local_tensor.copy_(local_shard) + return True + + def _dispatch_parameter( module: "nn.Module", name: str, @@ -627,6 +1095,8 @@ def _dispatch_parameter( if hasattr(orig_tensor, "device_mesh"): # dtensor device_mesh = getattr(orig_tensor, "device_mesh") placements = getattr(orig_tensor, "placements") + if _copy_into_existing_dtensor_shard(orig_tensor, tensor): + return if is_cuda_to_cuda: # Diagnostic: log expert stacked tensor dispatch (CUDA→CUDA DTensor path) logger.debug( @@ -844,12 +1314,23 @@ def all_ranks_load_weights( # Get checkpoint handler from model (delegates weight transforms to model-specific logic) # Pass the model's device so MoE expert stacking can happen on GPU (13x faster). - model_device = next(model.parameters()).device if any(True for _ in model.parameters()) else None + model_device = None + if hasattr(model, "parameters"): + try: + model_device = next(model.parameters()).device + except StopIteration: + model_device = None handler = None if hasattr(model, "get_checkpoint_handler"): ep_rank = _ps.ep_rank if _ps.ep_enabled else 0 ep_size = _ps.ep_size if _ps.ep_enabled else 1 checkpoint_keys = _get_checkpoint_keys(weights_path) + model_dtype = None + if hasattr(model, "parameters"): + try: + model_dtype = next(model.parameters()).dtype + except StopIteration: + model_dtype = None handler = model.get_checkpoint_handler( checkpoint_keys=checkpoint_keys or set(), ep_rank=ep_rank, @@ -857,6 +1338,7 @@ def all_ranks_load_weights( is_broadcast=False, weights_path=weights_path, device=model_device, + dtype=model_dtype, ) # Retry loading state dict on OSError (e.g., HuggingFace download issues) @@ -1070,19 +1552,26 @@ def rank0_load_and_broadcast_weights( _ps = get_parallel_state() global_rank = _ps.global_rank torch_device = torch.device(init_device) - # Get checkpoint handler from model. # For the broadcast path, is_broadcast=True tells the handler that rank 0 buffers # ALL experts (ep_size=1). EP slicing is handled later by ParallelPlan.shard_tensor(). handler = None if hasattr(model, "get_checkpoint_handler"): checkpoint_keys = _get_checkpoint_keys(weights_path) + model_dtype = None + if hasattr(model, "parameters"): + try: + model_dtype = next(model.parameters()).dtype + except StopIteration: + model_dtype = None handler = model.get_checkpoint_handler( checkpoint_keys=checkpoint_keys or set(), ep_rank=0, ep_size=1, is_broadcast=True, weights_path=weights_path, + device=torch_device, + dtype=model_dtype, ) skip_key_fn = handler.get_skip_key_fn() if handler is not None else None @@ -1297,6 +1786,403 @@ def rank0_load_and_broadcast_weights( post_process_after_weight_loading(model, buffer_dict, parameter_names_to_load, dtensor_factory) +@torch.no_grad() +def grouped_load_weights( + model: Union["nn.Module", "PreTrainedModel"], + weights_path: str, + init_device: Literal["cpu", "cuda", "npu"] = "cuda", + dtensor_factory: Optional[Callable[["torch.Tensor", Any, Any], "torch.Tensor"]] = None, +): + """Load weights with a hybrid dense/expert grouped strategy. + + Dense/shared tensors are loaded once per node and broadcast inside a + node-local load group. Expert tensors are loaded once per ``ep_fsdp`` group + leader and fanned out only to the ranks that share that expert shard. + """ + if not dist.is_available() or not dist.is_initialized(): + logger.warning_once("Distributed environment not initialized, falling back to all_ranks_load_weights.") + return all_ranks_load_weights(model, weights_path, init_device, dtensor_factory) + + _ps = get_parallel_state() + fanout_group = _get_grouped_weight_load_group(_ps) + if fanout_group is None: + logger.info_rank0("Grouped weight loading requires EP/FSDP groups; using rank-0 load fallback.") + return rank0_load_and_broadcast_weights(model, weights_path, init_device, dtensor_factory) + + buffer_dict = {name: buffer.clone() for name, buffer in model.named_buffers()} + parameter_names_to_load = {name for name, _ in model.named_parameters()} + + _compiled_key_map = _build_compiled_key_map(parameter_names_to_load, buffer_dict) + + _shrink_expert_params_for_ep(model) + model.to_empty(device=init_device) + + parallel_plan = None + if hasattr(model, "get_parallel_plan"): + parallel_plan = model.get_parallel_plan() + + fanout_ranks = dist.get_process_group_ranks(fanout_group) + fanout_src = fanout_ranks[0] + dense_group = _get_grouped_dense_weight_load_group() + if dense_group is None: + dense_ranks = [_ps.global_rank] + else: + dense_ranks = dist.get_process_group_ranks(dense_group) + dense_src = dense_ranks[0] + global_rank = _ps.global_rank + is_group_leader = global_rank == fanout_src + is_dense_leader = global_rank == dense_src + torch_device = torch.device(init_device) + model_device = None + model_dtype = None + if hasattr(model, "parameters"): + try: + first_param = next(model.parameters()) + model_device = first_param.device + model_dtype = first_param.dtype + except StopIteration: + model_device = None + model_dtype = None + checkpoint_keys = _get_checkpoint_keys(weights_path) if hasattr(model, "get_checkpoint_handler") else None + dense_handler = None + expert_handler = None + if hasattr(model, "get_checkpoint_handler"): + dense_handler = model.get_checkpoint_handler( + checkpoint_keys=checkpoint_keys or set(), + ep_rank=0, + ep_size=1, + is_broadcast=False, + weights_path=weights_path, + device=model_device, + dtype=model_dtype, + ) + if hasattr(model, "get_checkpoint_handler") and is_group_leader: + ep_rank = _ps.ep_rank if _ps.ep_enabled else 0 + ep_size = _ps.ep_size if _ps.ep_enabled else 1 + expert_handler = model.get_checkpoint_handler( + checkpoint_keys=checkpoint_keys or set(), + ep_rank=ep_rank, + ep_size=ep_size, + is_broadcast=False, + weights_path=weights_path, + device=model_device, + dtype=model_dtype, + ) + dense_skip_key_fn = dense_handler.get_skip_key_fn() if dense_handler is not None else None + expert_skip_key_fn = expert_handler.get_skip_key_fn() if expert_handler is not None else None + + if _ps.pp_enabled: + local_expected = parameter_names_to_load | set(buffer_dict.keys()) + all_expected = [None] * dist.get_world_size() + dist.all_gather_object(all_expected, local_expected) + global_expected_keys = set().union(*all_expected) + else: + global_expected_keys = None + + dense_state_dict_iterators = _load_state_dict(weights_path) + expert_state_dict_iterators = _load_state_dict(weights_path) + shard_count = len(dense_state_dict_iterators) + leader_count = dist.get_world_size() // len(fanout_ranks) + dense_leader_count = dist.get_world_size() // len(dense_ranks) + logger.info_rank0( + "grouped_load_weights: " + f"{shard_count=} " + f"dense_group_size={len(dense_ranks)} dense_leader_count={dense_leader_count} " + f"fanout_group_size={len(fanout_ranks)} leader_count={leader_count}" + ) + prefetch_count = _get_grouped_weight_load_prefetch_count(shard_count) + + _expected_skip_keys = set() + _expected_skip_prefixes = set() + for fqn, mod in model.named_modules(): + if getattr(mod, "_qlora_expected_skip_keys", None): + for suffix in mod._qlora_expected_skip_keys: + key = f"{fqn}.{suffix}" if fqn else suffix + _expected_skip_keys.add(key) + _expected_skip_prefixes.add(key) + + parameter_names_to_load -= _expected_skip_keys + + def _should_skip_qlora_expert_key(key: str, prefixes: set) -> bool: + for prefix in prefixes: + parts = prefix.rsplit(".", 1) + if len(parts) != 2: + continue + base, proj = parts + if key.startswith(base + ".") and f".{proj}." in key: + return True + return False + + def _dispatch_loaded_tensor(param_name: str, param_tensor: torch.Tensor, *, expect_expert: bool) -> None: + model_name = _compiled_key_map.get(param_name, param_name) + is_expert = _is_expert_parameter_name(model_name, parallel_plan) + if expect_expert != is_expert: + raise RuntimeError( + f"Grouped weight loading misrouted {'expert' if expect_expert else 'dense'} tensor " + f"{param_name} -> {model_name}" + ) + + if param_name in _expected_skip_keys or model_name in _expected_skip_keys: + return + if _expected_skip_prefixes and _should_skip_qlora_expert_key(model_name, _expected_skip_prefixes): + return + if model_name in buffer_dict: + buffer_dict[model_name] = param_tensor.clone() + return + if model_name in parameter_names_to_load: + parameter_names_to_load.discard(model_name) + _dispatch_parameter(model, model_name, param_tensor, dtensor_factory, parallel_plan) + return + if global_expected_keys is None or param_name not in global_expected_keys: + logger.warning_rank0(f"Unexpected key in state dict: {param_name}.") + + def _broadcast_queue_and_dispatch( + queue: List[Tuple[str, torch.Tensor]], + *, + group, + object_group=None, + src: int, + is_source: bool, + expect_expert: bool, + scatter_group_size: int = 1, + ) -> None: + # group=None here means "no replication group" — every rank loads + # locally and is its own source. Routing through the size-1 fast path + # below avoids issuing collectives on the default world group with + # disagreeing src. + if group is None: + group_size = 1 + else: + try: + group_size = dist.get_world_size(group) + except TypeError: + if hasattr(dist, "get_process_group_ranks"): + group_size = len(dist.get_process_group_ranks(group)) + else: + group_size = dist.get_world_size() + if group_size == 1: + local_batch = [] + if is_source: + for name, tensor in queue: + transfer_mode = "broadcast" + target_shape = tensor.shape + if expect_expert and scatter_group_size > 1: + model_name = _compiled_key_map.get(name, name) + if model_name in parameter_names_to_load: + maybe_target_shape = _get_expert_scatter_target_shape( + model, model_name, tensor, parallel_plan, _ps + ) + if maybe_target_shape is not None: + target_shape = torch.Size(maybe_target_shape) + transfer_mode = "expert_scatter" + local_batch.append((name, torch.Size(target_shape), tensor.dtype, transfer_mode)) + + for name, shape, dtype, transfer_mode in local_batch: + _, source_tensor = queue.pop(0) + 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) + _dispatch_loaded_tensor(name, tensor, expect_expert=expect_expert) + del tensor + del source_tensor + return + + if is_source: + batch_meta = [] + for name, tensor in queue: + transfer_mode = "broadcast" + target_shape = tensor.shape + if expect_expert and scatter_group_size > 1: + model_name = _compiled_key_map.get(name, name) + if model_name in parameter_names_to_load: + maybe_target_shape = _get_expert_scatter_target_shape( + model, model_name, tensor, parallel_plan, _ps + ) + if maybe_target_shape is not None: + target_shape = torch.Size(maybe_target_shape) + transfer_mode = "expert_scatter" + batch_meta.append((name, torch.Size(target_shape), tensor.dtype, transfer_mode)) + else: + batch_meta = None + + batch_meta = [batch_meta] + _broadcast_object_list(batch_meta, src=src, group=object_group if object_group is not None else group) + batch_meta = batch_meta[0] + + for name, shape, dtype, transfer_mode in batch_meta: + if is_source: + _, source_tensor = queue.pop(0) + if transfer_mode == "expert_scatter": + full_tensor, scatter_list = _build_group_scatter_list( + source_tensor, + tuple(shape), + scatter_group_size, + torch_device, + ) + tensor = scatter_list[0].clone() + 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) + full_tensor = None + scatter_list = None + else: + tensor = torch.empty(shape, dtype=dtype, device=torch_device) + source_tensor = None + full_tensor = None + scatter_list = None + + if transfer_mode == "expert_scatter": + dist.scatter(tensor, scatter_list=scatter_list, src=src, group=group) + else: + dist.broadcast(tensor, src=src, group=group) + _dispatch_loaded_tensor(name, tensor, expect_expert=expect_expert) + + del tensor + if is_source: + del source_tensor + if full_tensor is not None: + del full_tensor + if scatter_list is not None: + del scatter_list + + def _normalize_grouped_checkpoint_key(key: str) -> Optional[str]: + if _matches_checkpoint_skip_key_pattern(key, model): + return None + converted_key = _convert_weight_key(key, model) + if converted_key != key and _matches_checkpoint_skip_key_pattern(converted_key, model): + return None + return _normalize_checkpoint_key_for_filter(converted_key) + + def _should_skip_dense_key(key: str) -> bool: + normalized = _normalize_grouped_checkpoint_key(key) + if normalized is None or _is_checkpoint_expert_key(normalized): + return True + if dense_skip_key_fn is not None: + return dense_skip_key_fn(normalized) + return False + + def _should_skip_grouped_expert_key(key: str) -> bool: + normalized = _normalize_grouped_checkpoint_key(key) + if normalized is None or not _is_checkpoint_expert_key(normalized): + return True + if expert_skip_key_fn is not None: + return expert_skip_key_fn(normalized) + return False + + logger.info_rank0( + f"Grouped loading enabled: dense/shared tensors use dense_group_size={len(dense_ranks)} dense_src={dense_src}" + ) + dense_prefetched = None + if is_dense_leader: + dense_prefetched = _prefetch_shards_filtered( + dense_state_dict_iterators, + _should_skip_dense_key, + prefetch_count=1, + ) + + expert_prefetched = None + if is_group_leader: + logger.info_rank0( + f"Grouped loading enabled: expert tensors use fanout_group_size={len(fanout_ranks)} " + f"ep_rank={_ps.ep_rank} ep_size={_ps.ep_size} prefetch_count={prefetch_count}" + ) + expert_prefetched = _prefetch_shards_filtered( + expert_state_dict_iterators, + _should_skip_grouped_expert_key, + prefetch_count=prefetch_count, + ) + + shard_range = ( + tqdm( + range(shard_count), + desc="Loading checkpoint shards", + disable=global_rank != 0 or int(os.getenv("LOCAL_RANK", "-1")) > 0, + ) + if global_rank == 0 + else range(shard_count) + ) + + dense_queue: List[Tuple[str, torch.Tensor]] = [] + expert_queue: List[Tuple[str, torch.Tensor]] = [] + + for _ in shard_range: + if is_dense_leader: + state_dict, _skipped = next(dense_prefetched) + for key, tensor in state_dict.items(): + key = _convert_weight_key(key, model) + results = dense_handler.on_load_weight(key, tensor) if dense_handler is not None else [(key, tensor)] + for result_name, result_tensor in results: + dense_queue.append((result_name, result_tensor)) + del state_dict + _broadcast_queue_and_dispatch( + dense_queue, + group=dense_group, + src=dense_src, + is_source=is_dense_leader, + expect_expert=False, + ) + + if is_group_leader: + state_dict, skipped_keys = next(expert_prefetched) + for skipped_key in skipped_keys: + normalized = _normalize_grouped_checkpoint_key(skipped_key) + if normalized is None or not _is_checkpoint_expert_key(normalized): + continue + if expert_skip_key_fn is not None and expert_skip_key_fn(normalized): + expert_queue.extend(expert_handler.on_skip_weight(normalized)) + for key, tensor in state_dict.items(): + key = _convert_weight_key(key, model) + results = expert_handler.on_load_weight(key, tensor) if expert_handler is not None else [(key, tensor)] + for result_name, result_tensor in results: + expert_queue.append((result_name, result_tensor)) + del state_dict + + _broadcast_queue_and_dispatch( + expert_queue, + group=fanout_group, + src=fanout_src, + is_source=is_group_leader, + expect_expert=True, + scatter_group_size=len(fanout_ranks), + ) + + empty_cache() + + if dense_handler is not None: + if is_dense_leader: + dense_queue.extend(dense_handler.on_load_complete()) + _broadcast_queue_and_dispatch( + dense_queue, + group=dense_group, + src=dense_src, + is_source=is_dense_leader, + expect_expert=False, + ) + + if expert_handler is not None and is_group_leader: + expert_queue.extend(expert_handler.on_load_complete()) + _broadcast_queue_and_dispatch( + expert_queue, + group=fanout_group, + src=fanout_src, + is_source=is_group_leader, + expect_expert=True, + scatter_group_size=len(fanout_ranks), + ) + + post_process_after_weight_loading( + model, + buffer_dict, + parameter_names_to_load, + dtensor_factory, + qlora_skip_prefixes=_expected_skip_prefixes, + qlora_skip_fn=_should_skip_qlora_expert_key, + ) + + def post_process_after_weight_loading( model: Union["nn.Module", "PreTrainedModel"], buffer_dict, @@ -1420,6 +2306,104 @@ def _save_state_dict( torch.save(state_dict, path_to_save) +def _distributed_barrier_on_current_device() -> None: + if not dist.is_available() or not dist.is_initialized(): + return + + barrier_kwargs = {} + if dist.get_backend() == "nccl": + barrier_kwargs["device_ids"] = [get_device_id()] + dist.barrier(**barrier_kwargs) + + +def _get_save_sync_group(): + """Get or create a dedicated process group for checkpoint-save coordination.""" + global _save_sync_group, _save_sync_group_backend + if not dist.is_initialized(): + return None + + backend = os.getenv("XORL_SAVE_SYNC_BACKEND", "filesystem").strip().lower() + if backend in {"", "default", "filesystem", dist.get_backend()}: + return None + + if _save_sync_group is None or _save_sync_group_backend != backend: + timeout_sec = int(os.getenv("XORL_SAVE_SYNC_TIMEOUT_SEC", "7200")) + _save_sync_group = dist.new_group(backend=backend, timeout=timedelta(seconds=timeout_sec)) + _save_sync_group_backend = backend + return _save_sync_group + + +def _filesystem_save_barrier(output_dir: Union[str, "os.PathLike"], barrier_name: str) -> None: + timeout_sec = int(os.getenv("XORL_SAVE_SYNC_TIMEOUT_SEC", "7200")) + poll_sec = float(os.getenv("XORL_SAVE_SYNC_POLL_SEC", "1.0")) + run_id_raw = ( + os.getenv("TORCHELASTIC_RUN_ID") + or os.getenv("XORL_SAVE_SYNC_ID") + or f"{os.getenv('MASTER_ADDR', 'local')}_{os.getenv('MASTER_PORT', '0')}" + ) + run_id = re.sub(r"[^A-Za-z0-9_.-]", "_", run_id_raw) + barrier_root = os.path.join(output_dir, ".xorl_save_barriers", run_id) + os.makedirs(barrier_root, exist_ok=True) + + rank = dist.get_rank() + world_size = dist.get_world_size() + arrival_marker = os.path.join(barrier_root, f"{barrier_name}.rank{rank:05d}") + done_marker = os.path.join(barrier_root, f"{barrier_name}.done") + + with open(arrival_marker, "w", encoding="utf-8") as f: + f.write("\n") + + deadline = time.monotonic() + timeout_sec + arrival_prefix = f"{barrier_name}.rank" + while True: + arrival_count = sum(1 for entry in os.scandir(barrier_root) if entry.name.startswith(arrival_prefix)) + if arrival_count >= world_size: + break + if time.monotonic() >= deadline: + raise TimeoutError( + f"Timed out waiting for filesystem save barrier {barrier_name}: " + f"{arrival_count}/{world_size} ranks arrived." + ) + time.sleep(poll_sec) + + if rank == 0 and not os.path.exists(done_marker): + with open(done_marker, "w", encoding="utf-8") as f: + f.write("\n") + + while not os.path.exists(done_marker): + if time.monotonic() >= deadline: + raise TimeoutError(f"Timed out waiting for filesystem save barrier completion: {barrier_name}.") + time.sleep(poll_sec) + + +def _distributed_save_barrier( + output_dir: Optional[Union[str, "os.PathLike"]] = None, + barrier_name: str = "save-sync", +) -> None: + """Barrier helper for checkpoint save paths. + + Saving mostly coordinates CPU/filesystem work, and some clusters have shown + NCCL barrier instability once the training/load collectives are done. Use a + dedicated save-sync process group when configured, falling back to the + current-device barrier only when the backends already match. + """ + if not dist.is_available() or not dist.is_initialized(): + return + + backend = os.getenv("XORL_SAVE_SYNC_BACKEND", "filesystem").strip().lower() + if backend == "filesystem": + if output_dir is None: + raise ValueError("output_dir is required for filesystem save barriers.") + _filesystem_save_barrier(output_dir, barrier_name) + return + + group = _get_save_sync_group() + if group is None: + _distributed_barrier_on_current_device() + else: + dist.barrier(group=group) + + @torch.no_grad() def save_model_weights( output_dir: Union[str, "os.PathLike"], @@ -1469,7 +2453,7 @@ def save_model_weights( empty_cache() if global_rank is not None and dist.is_initialized(): # avoid process hanging synchronize() - dist.barrier() + _distributed_barrier_on_current_device() if global_rank is None or global_rank == 0: full_state_dict[name] = tensor.detach().cpu() @@ -1503,9 +2487,149 @@ 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 +def save_model_weights_distributed( + output_dir: Union[str, "os.PathLike"], + state_dict: Dict[str, "torch.Tensor"], + save_dtype: Optional[Union[str, "torch.dtype"]] = "bfloat16", + shard_size: int = 5_000_000_000, + safe_serialization: bool = True, + model_assets: Optional[Sequence["ModelAssets"]] = None, +) -> None: + """Save a full model checkpoint with one writer per node. + This helper is intended for large EP/FSDP models where every rank must + participate in DTensor materialization, but rank-0-only writing would create + a single-node CPU/I/O bottleneck. We deterministically assign output shards + to the local-rank-0 writer on each node while every rank still participates + in the required ``full_tensor()`` collectives. + """ + if not safe_serialization: + raise ValueError("Distributed weight saving currently requires safe_serialization=True.") + + if not dist.is_available() or not dist.is_initialized() or dist.get_world_size() <= 1: + global_rank = dist.get_rank() if dist.is_available() and dist.is_initialized() else None + save_model_weights( + output_dir=output_dir, + state_dict=state_dict, + global_rank=global_rank, + save_dtype=save_dtype, + shard_size=shard_size, + safe_serialization=safe_serialization, + model_assets=model_assets, + ) + return + + global_rank = dist.get_rank() + local_world_size = max(1, int(os.environ.get("LOCAL_WORLD_SIZE", "1"))) + local_rank = int(os.environ.get("LOCAL_RANK", str(global_rank % local_world_size))) + writer_ranks = list(range(0, dist.get_world_size(), local_world_size)) + is_writer = local_rank == 0 + + if is_writer: + os.makedirs(output_dir, exist_ok=True) + else: + os.makedirs(output_dir, exist_ok=True) + + _distributed_save_barrier(output_dir, "materialize-start") + + is_sharded, total_size, weight_map = _get_shard_info(state_dict, save_dtype, shard_size, safe_serialization) + shard_to_keys: "OrderedDict[str, List[str]]" = OrderedDict() + for name, file_name in weight_map.items(): + shard_to_keys.setdefault(file_name, []).append(name) + + shard_to_writer = { + file_name: writer_ranks[idx % len(writer_ranks)] for idx, file_name in enumerate(shard_to_keys.keys()) + } + total_tensors = len(state_dict) + total_shards = len(shard_to_keys) + + dtype_target = getattr(torch, save_dtype) if isinstance(save_dtype, str) else save_dtype + materialized_count = 0 + last_progress_log = time.monotonic() + progress_log_interval = float(os.getenv("XORL_SAVE_PROGRESS_LOG_INTERVAL_SEC", "15")) + + if global_rank == 0: + logger.info( + "Distributed save starting: " + f"{total_shards} shards, {total_tensors} tensors, " + f"{len(writer_ranks)} writer ranks" + ) + + for shard_idx, (file_name, shard_keys) in enumerate(shard_to_keys.items(), start=1): + owner_rank = shard_to_writer[file_name] + keep_shard = owner_rank == global_rank + shard_state_dict: "OrderedDict[str, torch.Tensor]" = OrderedDict() if keep_shard else OrderedDict() + + for shard_tensor_idx, name in enumerate(shard_keys, start=1): + tensor = state_dict[name] + materialized = _materialize_tensor_for_save(tensor, dst_rank=owner_rank) + if keep_shard and materialized is not None: + cpu_tensor = materialized.detach().cpu() + if dtype_target is not None and cpu_tensor.dtype != dtype_target and cpu_tensor.is_floating_point(): + cpu_tensor = cpu_tensor.to(dtype=dtype_target) + shard_state_dict[name] = cpu_tensor + + del materialized + materialized_count += 1 + if materialized_count % 32 == 0: + empty_cache() + + now = time.monotonic() + if global_rank == 0 and now - last_progress_log >= progress_log_interval: + logger.info( + "Distributed save progress: " + f"shard {shard_idx}/{total_shards} " + f"({shard_tensor_idx}/{len(shard_keys)} tensors in current shard), " + f"materialized {materialized_count}/{total_tensors} tensors, " + f"current owner_rank={owner_rank}, file={file_name}" + ) + last_progress_log = now + + if keep_shard: + _save_state_dict(shard_state_dict, os.path.join(output_dir, file_name), safe_serialization) + logger.info( + f"Distributed save wrote shard {file_name} on rank {global_rank} " + f"({len(shard_keys)} tensors, shard {shard_idx}/{total_shards})" + ) + shard_state_dict.clear() + + if global_rank == 0 and (shard_idx == 1 or shard_idx == total_shards): + logger.info( + "Distributed save progress: " + f"materialized shard {shard_idx}/{total_shards} " + f"({materialized_count}/{total_tensors} tensors); " + f"current owner_rank={owner_rank}, file={file_name}" + ) + + empty_cache() + _distributed_save_barrier(output_dir, "shards-written") + + if global_rank == 0: + if is_sharded: + index = { + "metadata": {"total_size": total_size}, + "weight_map": weight_map, + } + index_file = SAFE_WEIGHTS_INDEX_NAME if safe_serialization else WEIGHTS_INDEX_NAME + with open(os.path.join(output_dir, index_file), "w", encoding="utf-8") as f: + content = json.dumps(index, indent=2, sort_keys=True) + "\n" + f.write(content) + logger.info(f"Distributed model weight splits saved in {output_dir}.") + else: + only_file = next(iter(shard_to_keys)) + logger.info(f"Distributed model weights saved at {os.path.join(output_dir, only_file)}.") + + if model_assets is not None: + for model_asset in model_assets: + if hasattr(model_asset, "save_pretrained"): + model_asset.save_pretrained(output_dir) + else: + logger.warning(f"Model asset {model_asset} should implement `save_pretrained`.") + + _distributed_save_barrier(output_dir, "index-written") + + +def save_model_assets(output_dir: Union[str, "os.PathLike"], model_assets: Sequence["ModelAssets"]): for model_asset in model_assets: if hasattr(model_asset, "save_pretrained"): try: @@ -1567,6 +2691,14 @@ def compute_loss( return loss_fn(**loss_kwargs) +GradientCheckpointingMethod = Literal[ + "recompute_full_layer", + "recompute_before_dispatch", + "no_recompute", +] +DEFAULT_GRADIENT_CHECKPOINTING_METHOD: GradientCheckpointingMethod = "recompute_full_layer" + + class GradientCheckpointingLayer(nn.Module): """Base class for layers with gradient checkpointing. @@ -1591,12 +2723,13 @@ class GradientCheckpointingLayer(nn.Module): """ gradient_checkpointing = False + _gradient_checkpointing_method: GradientCheckpointingMethod = DEFAULT_GRADIENT_CHECKPOINTING_METHOD def __call__(self, *args, **kwargs): if ( self.gradient_checkpointing and self.training - and getattr(self, "_gradient_checkpointing_method", "recompute_full_layer") == "recompute_full_layer" + and self._gradient_checkpointing_method == DEFAULT_GRADIENT_CHECKPOINTING_METHOD ): return self._gradient_checkpointing_func(partial(super().__call__, **kwargs), *args) return super().__call__(*args, **kwargs) @@ -1621,6 +2754,7 @@ class MoEGradientCheckpointingLayer(nn.Module): """ gradient_checkpointing = False + _gradient_checkpointing_method: GradientCheckpointingMethod = DEFAULT_GRADIENT_CHECKPOINTING_METHOD def _pre_mlp_forward(self, hidden_states, **kwargs): """Layernorm → attention → layernorm. Override per model. @@ -1654,13 +2788,13 @@ def _moe_forward(self, hidden_states, output_router_logits=False, **kwargs): 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 + from xorl.distributed.moe.deepep import sync_pending_combine # noqa: PLC0415 + from xorl.models.layers.moe.moe_block import MoEBlock # noqa: PLC0415 _selective = ( self.training - and getattr(self, "gradient_checkpointing", False) - and getattr(self, "_gradient_checkpointing_method", "recompute_full_layer") != "recompute_full_layer" + and self.gradient_checkpointing + and self._gradient_checkpointing_method != DEFAULT_GRADIENT_CHECKPOINTING_METHOD ) _is_moe = isinstance(self.mlp, MoEBlock) diff --git a/src/xorl/models/registry.py b/src/xorl/models/registry.py index 4697302e..a3772d72 100644 --- a/src/xorl/models/registry.py +++ b/src/xorl/models/registry.py @@ -21,6 +21,7 @@ class _ModelRegistry: # Keyed by model_arch modeling_path: List[str] = field(default_factory=list) model_arch_name_to_cls: Dict[str, Union[Type[nn.Module], str]] = field(default_factory=dict) + import_errors: Dict[str, Exception] = field(default_factory=dict) def __post_init__(self): for modeling_path in self.modeling_path: @@ -46,7 +47,8 @@ def _mapping_model_arch_name_to_cls(self, modeling_path: str): try: module = importlib.import_module(name) except Exception as e: - logger.warning(f"Ignore import error when loading {name}. {e}") + logger.warning(f"Import error when loading {name}: {e}") + self.import_errors[name] = e continue if hasattr(module, "ModelClass"): entry = module.ModelClass diff --git a/src/xorl/models/transformers/__init__.py b/src/xorl/models/transformers/__init__.py index a4bec082..0edc5da7 100644 --- a/src/xorl/models/transformers/__init__.py +++ b/src/xorl/models/transformers/__init__.py @@ -1,4 +1,6 @@ from . import ( + deepseek_v3, + glm4_moe, llama3, qwen2, qwen3, @@ -9,10 +11,12 @@ __all__ = [ + "deepseek_v3", + "glm4_moe", "llama3", "qwen2", "qwen3", "qwen3_5", - "qwen3_moe", "qwen3_5_moe", + "qwen3_moe", ] diff --git a/src/xorl/models/transformers/deepseek_v3/__init__.py b/src/xorl/models/transformers/deepseek_v3/__init__.py new file mode 100644 index 00000000..8dfbe7b4 --- /dev/null +++ b/src/xorl/models/transformers/deepseek_v3/__init__.py @@ -0,0 +1,14 @@ +from .configuration_deepseek_v3 import DeepseekV3Config +from .modeling_deepseek_v3 import ( + DeepseekV3ForCausalLM, + DeepseekV3Model, + DeepseekV3PreTrainedModel, +) + + +__all__ = [ + "DeepseekV3Config", + "DeepseekV3Model", + "DeepseekV3ForCausalLM", + "DeepseekV3PreTrainedModel", +] diff --git a/src/xorl/models/transformers/deepseek_v3/checkpoint_handler.py b/src/xorl/models/transformers/deepseek_v3/checkpoint_handler.py new file mode 100644 index 00000000..5b1d8c5f --- /dev/null +++ b/src/xorl/models/transformers/deepseek_v3/checkpoint_handler.py @@ -0,0 +1,370 @@ +"""Checkpoint handler for DeepseekV3 / Kimi-K2.5.""" + +import math +import re +import warnings +from collections import defaultdict +from typing import Callable, Dict, List, Optional, Set, Tuple + +import torch + +from ...checkpoint_handlers.base import CheckpointHandler +from ...checkpoint_handlers.buffers import ( + ExpertWeightBuffer, + parse_expert_key, +) + + +class DeepseekV3CheckpointHandler(CheckpointHandler): + _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$") + _COMPRESSED_EXPERT_PATTERN = re.compile( + r"^model\.layers\.(\d+)\.mlp\.experts\.(\d+)\.(gate|up|down)_proj\.(weight_packed|weight_scale|weight_shape)$" + ) + _COMPRESSED_EXPERT_SUFFIXES = frozenset({"weight_packed", "weight_scale", "weight_shape"}) + + def __init__( + self, + num_experts: int, + ep_rank: int = 0, + ep_size: int = 1, + checkpoint_has_per_expert: bool = True, + packed_expert_num_bits: int = 4, + packed_expert_group_size: int = 32, + device: Optional[torch.device] = None, + dtype: Optional[torch.dtype] = None, + ): + self._dequant_device = None if device is None or device.type == "meta" else device + self._output_dtype = dtype + self._expert_buffer: Optional[ExpertWeightBuffer] = None + if checkpoint_has_per_expert: + self._expert_buffer = ExpertWeightBuffer( + num_experts, + ep_rank=ep_rank, + ep_size=ep_size, + device=self._dequant_device, + ) + 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]] = {} + self._packed_expert_num_bits = packed_expert_num_bits + self._packed_expert_group_size = packed_expert_group_size + self._packed_expert_pending: Dict[Tuple[int, int, str], Dict[str, torch.Tensor]] = {} + self._skipped_packed_expert_suffixes: Dict[Tuple[int, int, str], Set[str]] = defaultdict(set) + + def _normalize_key(self, key: str) -> Optional[str]: + if key.startswith("vision_tower.") or key.startswith("mm_projector."): + return None + if key.startswith("language_model."): + return key.removeprefix("language_model.") + return key + + 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 _parse_compressed_expert_key(self, key: str) -> Optional[Tuple[int, int, str, str]]: + match = self._COMPRESSED_EXPERT_PATTERN.match(key) + if match is None: + return None + return int(match.group(1)), int(match.group(2)), match.group(3), match.group(4) + + def _unpack_packed_int32_tensor(self, tensor: torch.Tensor, original_shape: Tuple[int, int]) -> torch.Tensor: + if tensor.dtype != torch.int32: + tensor = tensor.to(torch.int32) + if tensor.ndim != 2: + raise ValueError(f"Expected packed expert tensor to be rank-2, got shape {tuple(tensor.shape)}") + + num_bits = self._packed_expert_num_bits + pack_factor = 32 // num_bits + mask = (1 << num_bits) - 1 + rows, cols = original_shape + + bit_shifts = torch.arange(pack_factor, device=tensor.device, dtype=torch.int32) * num_bits + unpacked = ((tensor.unsqueeze(-1) >> bit_shifts) & mask).reshape(rows, -1) + unpacked = unpacked[:, :cols] + + offset = 1 << (num_bits - 1) + return (unpacked - offset).to(torch.int8) + + def _dequantize_packed_expert_weight( + self, + packed: torch.Tensor, + scale: torch.Tensor, + shape: torch.Tensor, + ) -> torch.Tensor: + if self._dequant_device is not None: + packed = packed.to(self._dequant_device) + scale = scale.to(self._dequant_device) + shape = shape.to(self._dequant_device) + + original_shape = tuple(int(dim) for dim in shape.flatten().tolist()) + if len(original_shape) != 2: + raise ValueError(f"Expected rank-2 expert weight shape, got {original_shape}") + + rows, cols = original_shape + groups = math.ceil(cols / self._packed_expert_group_size) + unpacked = self._unpack_packed_int32_tensor(packed, (rows, cols)) + compute_dtype = scale.dtype if scale.is_floating_point() else torch.float32 + dequantized = unpacked.to(compute_dtype) + + if scale.ndim == 0: + result = dequantized * scale.to(compute_dtype) + return result.to(self._output_dtype) if self._output_dtype is not None else result + + if scale.ndim == 1: + if scale.numel() == rows: + result = dequantized * scale.to(compute_dtype).view(rows, 1) + return result.to(self._output_dtype) if self._output_dtype is not None else result + if scale.numel() == groups: + scale = scale.to(compute_dtype).view(1, groups) + else: + raise ValueError( + f"Unsupported packed expert scale shape {tuple(scale.shape)} for original shape {original_shape}" + ) + elif scale.ndim == 2: + if scale.shape == (rows, 1): + result = dequantized * scale.to(compute_dtype) + return result.to(self._output_dtype) if self._output_dtype is not None else result + if scale.shape[0] in (1, rows) and scale.shape[1] == groups: + scale = scale.to(compute_dtype) + elif scale.shape[1] in (1, rows) and scale.shape[0] == groups: + scale = scale.t().contiguous().to(compute_dtype) + else: + raise ValueError( + f"Unsupported packed expert scale shape {tuple(scale.shape)} for original shape {original_shape}" + ) + else: + raise ValueError(f"Unsupported packed expert scale rank {scale.ndim} for shape {original_shape}") + + padded_cols = groups * self._packed_expert_group_size + if padded_cols > cols: + dequantized = torch.nn.functional.pad(dequantized, (0, padded_cols - cols)) + dequantized = dequantized.unflatten(1, (groups, self._packed_expert_group_size)) + + if scale.shape[0] == 1 and rows != 1: + scale = scale.expand(rows, -1) + + result = (dequantized * scale.unsqueeze(-1)).flatten(1)[:, :cols].contiguous() + if self._output_dtype is not None and result.dtype != self._output_dtype: + result = result.to(self._output_dtype) + return result + + def _handle_compressed_expert_weight( + self, + key: str, + tensor: torch.Tensor, + ) -> Optional[List[Tuple[str, torch.Tensor]]]: + parsed = self._parse_compressed_expert_key(key) + if parsed is None: + return None + + layer_idx, expert_idx, proj, suffix = parsed + buffer_key = (layer_idx, expert_idx, proj) + pending = self._packed_expert_pending.setdefault(buffer_key, {}) + pending[suffix] = tensor + + if not self._COMPRESSED_EXPERT_SUFFIXES.issubset(pending): + return [] + + dense_weight = self._dequantize_packed_expert_weight( + packed=pending["weight_packed"], + scale=pending["weight_scale"], + shape=pending["weight_shape"], + ) + del self._packed_expert_pending[buffer_key] + + if self._expert_buffer is None: + return [(f"model.layers.{layer_idx}.mlp.experts.{expert_idx}.{proj}_proj.weight", dense_weight)] + + self._expert_buffer.add(layer_idx, expert_idx, proj, dense_weight) + return self._maybe_finalize_per_expert_merge(layer_idx, proj) + + 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: + # xorl checkpoints already store fused expert weights in native + # [experts, hidden, 2 * intermediate] layout. Only EP slicing is needed. + return [(key, self._slice_expert_tensor_for_ep(tensor))] + + down_match = self._FUSED_EXPERT_DOWN_PATTERN.match(key) + if down_match is not None: + # xorl checkpoints already store fused expert weights in native + # [experts, intermediate, hidden] layout. Only EP slicing is needed. + return [(key, self._slice_expert_tensor_for_ep(tensor))] + + split_match = self._STACKED_EXPERT_SPLIT_PATTERN.match(key) + if split_match is None: + return None + + layer_idx = int(split_match.group(1)) + proj = split_match.group(2) + tensor = self._slice_expert_tensor_for_ep(tensor) + + if proj == "down": + return [(f"model.layers.{layer_idx}.mlp.experts.down_proj", tensor.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]]: + 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: + return None + + ep_start = self._expert_buffer.expert_start + ep_end = self._expert_buffer.expert_end + + def _should_skip(key: str) -> bool: + normalized_key = self._normalize_key(key) + if normalized_key is None: + return True + + compressed = self._parse_compressed_expert_key(normalized_key) + if compressed is not None: + _, expert_idx, _, _ = compressed + return expert_idx < ep_start or expert_idx >= ep_end + + parsed = parse_expert_key(normalized_key) + if parsed is None: + return False + _, expert_idx, _ = parsed + return expert_idx < ep_start or expert_idx >= ep_end + + return _should_skip + + def on_load_weight(self, key: str, tensor: torch.Tensor) -> List[Tuple[str, torch.Tensor]]: + key = self._normalize_key(key) + if key is None: + return [] + + if key.endswith(".input_scale"): + return [] + + compressed_expert_results = self._handle_compressed_expert_weight(key, tensor) + if compressed_expert_results is not None: + return compressed_expert_results + + stacked_expert_results = self._handle_stacked_expert_weights(key, tensor) + if stacked_expert_results is not None: + return stacked_expert_results + + 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) + return self._maybe_finalize_per_expert_merge(layer_idx, proj) + + return [(key, tensor)] + + def on_skip_weight(self, key: str) -> List[Tuple[str, torch.Tensor]]: + if self._expert_buffer is None: + return [] + key = self._normalize_key(key) + if key is None: + return [] + + compressed = self._parse_compressed_expert_key(key) + if compressed is not None: + layer_idx, expert_idx, proj, suffix = compressed + skipped = self._skipped_packed_expert_suffixes[(layer_idx, expert_idx, proj)] + skipped.add(suffix) + if self._COMPRESSED_EXPERT_SUFFIXES.issubset(skipped): + del self._skipped_packed_expert_suffixes[(layer_idx, expert_idx, proj)] + self._expert_buffer.count_skipped(layer_idx, proj) + return self._maybe_finalize_per_expert_merge(layer_idx, proj) + return [] + + parsed = parse_expert_key(key) + if parsed is None: + return [] + layer_idx, _expert_idx, proj = parsed + self._expert_buffer.count_skipped(layer_idx, proj) + return self._maybe_finalize_per_expert_merge(layer_idx, proj) + + def on_load_complete(self) -> List[Tuple[str, torch.Tensor]]: + 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._packed_expert_pending: + pending_triplets = {key: sorted(value.keys()) for key, value in self._packed_expert_pending.items()} + warnings.warn(f"Incomplete packed expert triplets after loading: {pending_triplets}") + if self._skipped_packed_expert_suffixes: + pending_skips = {key: sorted(value) for key, value in self._skipped_packed_expert_suffixes.items()} + warnings.warn(f"Incomplete skipped packed expert triplets after loading: {pending_skips}") + return [] + + def on_save_weight(self, param_name: str, tensor: torch.Tensor) -> List[Tuple[str, torch.Tensor]]: + 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]) + ] + + return [(param_name, tensor)] diff --git a/src/xorl/models/transformers/deepseek_v3/configuration_deepseek_v3.py b/src/xorl/models/transformers/deepseek_v3/configuration_deepseek_v3.py new file mode 100644 index 00000000..79a0d782 --- /dev/null +++ b/src/xorl/models/transformers/deepseek_v3/configuration_deepseek_v3.py @@ -0,0 +1,235 @@ +"""DeepseekV3 / Kimi-K2.5 text-only configuration.""" + +from transformers.configuration_utils import PretrainedConfig + +from xorl.models.layers import rope_config_validation + + +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 DeepseekV3Config(PretrainedConfig): + model_type = "deepseek_v3" + + base_model_tp_plan = {} + base_model_pp_plan = { + "embed_tokens": (["input_ids"], ["inputs_embeds"]), + "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), + "norm": (["hidden_states"], ["hidden_states"]), + } + attribute_map = { + "num_local_experts": "n_routed_experts", + } + + def __init__( + self, + vocab_size=163840, + hidden_size=7168, + intermediate_size=18432, + moe_intermediate_size=2048, + num_hidden_layers=61, + num_attention_heads=64, + num_key_value_heads=64, + n_shared_experts=1, + n_routed_experts=384, + routed_scaling_factor=2.827, + kv_lora_rank=512, + q_lora_rank=1536, + qk_rope_head_dim=64, + v_head_dim=128, + qk_nope_head_dim=128, + n_group=1, + topk_group=1, + num_experts_per_tok=8, + first_k_dense_replace=1, + norm_topk_prob=True, + hidden_act="silu", + max_position_embeddings=262144, + initializer_range=0.02, + rms_norm_eps=1e-5, + use_cache=False, + pad_token_id=None, + bos_token_id=None, + eos_token_id=None, + tie_word_embeddings=False, + rope_theta=50000.0, + rope_scaling=None, + rope_interleave=True, + attention_bias=False, + attention_dropout=0.0, + output_router_logits=False, + router_aux_loss_coef=0.0, + topk_method="noaux_tc", + scoring_func="sigmoid", + _moe_implementation="eager", + **kwargs, + ): + self.vocab_size = vocab_size + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.moe_intermediate_size = moe_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.n_shared_experts = n_shared_experts + self.n_routed_experts = n_routed_experts + self.num_experts = n_routed_experts + self.routed_scaling_factor = routed_scaling_factor + self.kv_lora_rank = kv_lora_rank + self.q_lora_rank = q_lora_rank + self.qk_rope_head_dim = qk_rope_head_dim + self.v_head_dim = v_head_dim + self.qk_nope_head_dim = qk_nope_head_dim + self.qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim + self.head_dim = self.qk_head_dim + self.partial_rotary_factor = self.qk_rope_head_dim / self.qk_head_dim + self.n_group = n_group + self.topk_group = topk_group + self.num_experts_per_tok = num_experts_per_tok + self.first_k_dense_replace = first_k_dense_replace + self.norm_topk_prob = norm_topk_prob + self.hidden_act = hidden_act + self.max_position_embeddings = max_position_embeddings + self.initializer_range = initializer_range + self.rms_norm_eps = rms_norm_eps + self.use_cache = use_cache + self.pad_token_id = pad_token_id + self.bos_token_id = bos_token_id + self.eos_token_id = eos_token_id + self.tie_word_embeddings = tie_word_embeddings + self.rope_theta = rope_theta + self.rope_interleave = rope_interleave + self.attention_bias = attention_bias + self.attention_dropout = attention_dropout + self.output_router_logits = output_router_logits + self.router_aux_loss_coef = router_aux_loss_coef + self.topk_method = topk_method + self.scoring_func = scoring_func + self._moe_implementation = _moe_implementation + self._rope_scaling = _cfg_to_dict(rope_scaling) + + 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) + + 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, + ) + + @property + def rope_scaling(self): + return self._rope_scaling + + @rope_scaling.setter + def rope_scaling(self, value): + self._rope_scaling = _cfg_to_dict(value) + + @property + def rope_parameters(self): + rope_params = { + "rope_type": "default", + "rope_theta": self.rope_theta, + "partial_rotary_factor": self.partial_rotary_factor, + } + 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") + 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) and "rope_theta" in value_dict: + self.rope_theta = value_dict["rope_theta"] + self._rope_scaling = value_dict + + @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) + + router_aux_loss_coef = getattr(text_config, "router_aux_loss_coef", None) + if router_aux_loss_coef is None: + router_aux_loss_coef = getattr(text_config, "aux_loss_alpha", 0.0) + + output_router_logits = getattr(text_config, "output_router_logits", None) + if output_router_logits is None: + output_router_logits = router_aux_loss_coef > 0 + + qk_nope_head_dim = getattr(text_config, "qk_nope_head_dim") + qk_rope_head_dim = getattr(text_config, "qk_rope_head_dim") + + return cls( + vocab_size=getattr(text_config, "vocab_size", getattr(hf_config, "vocab_size", 163840)), + hidden_size=getattr(text_config, "hidden_size"), + intermediate_size=getattr(text_config, "intermediate_size"), + moe_intermediate_size=getattr(text_config, "moe_intermediate_size"), + num_hidden_layers=getattr(text_config, "num_hidden_layers"), + num_attention_heads=getattr(text_config, "num_attention_heads"), + num_key_value_heads=getattr( + text_config, + "num_key_value_heads", + getattr(text_config, "num_attention_heads"), + ), + n_shared_experts=getattr(text_config, "n_shared_experts", 1), + n_routed_experts=getattr(text_config, "n_routed_experts"), + routed_scaling_factor=getattr(text_config, "routed_scaling_factor", 1.0), + kv_lora_rank=getattr(text_config, "kv_lora_rank"), + q_lora_rank=getattr(text_config, "q_lora_rank", None), + qk_rope_head_dim=qk_rope_head_dim, + v_head_dim=getattr(text_config, "v_head_dim"), + qk_nope_head_dim=qk_nope_head_dim, + n_group=getattr(text_config, "n_group", 1), + topk_group=getattr(text_config, "topk_group", 1), + num_experts_per_tok=getattr(text_config, "num_experts_per_tok"), + first_k_dense_replace=getattr(text_config, "first_k_dense_replace", 0), + norm_topk_prob=getattr(text_config, "norm_topk_prob", True), + 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-5), + use_cache=getattr(text_config, "use_cache", False), + 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)), + 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), + rope_interleave=getattr(text_config, "rope_interleave", True), + attention_bias=getattr(text_config, "attention_bias", False), + attention_dropout=getattr(text_config, "attention_dropout", 0.0), + output_router_logits=output_router_logits, + router_aux_loss_coef=router_aux_loss_coef, + topk_method=getattr(text_config, "topk_method", "noaux_tc"), + scoring_func=getattr(text_config, "scoring_func", "sigmoid"), + architectures=list(getattr(text_config, "architectures", ["DeepseekV3ForCausalLM"])), + ) + + +__all__ = ["DeepseekV3Config"] diff --git a/src/xorl/models/transformers/deepseek_v3/modeling_deepseek_v3.py b/src/xorl/models/transformers/deepseek_v3/modeling_deepseek_v3.py new file mode 100644 index 00000000..3850356c --- /dev/null +++ b/src/xorl/models/transformers/deepseek_v3/modeling_deepseek_v3.py @@ -0,0 +1,652 @@ +from typing import Optional, Tuple + +import torch +import torch.nn.functional as F +from torch import nn + +from xorl.distributed.moe.deepep import sync_pending_combine +from xorl.distributed.parallel_state import get_parallel_state +from xorl.distributed.sequence_parallel.strategy import get_cp_strategy +from xorl.models.base import XorlPreTrainedModel +from xorl.models.layers import ACT2FN, RMSNorm, RotaryEmbedding +from xorl.models.layers.attention import is_flash_attention, 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.layers.moe.routing_replay import get_replay_stage +from xorl.models.outputs import MoeCausalLMOutput, MoeModelOutput +from xorl.models.transformers.deepseek_v3 import parallelize +from xorl.models.transformers.deepseek_v3.checkpoint_handler import DeepseekV3CheckpointHandler +from xorl.models.transformers.deepseek_v3.configuration_deepseek_v3 import DeepseekV3Config +from xorl.models.transformers.deepseek_v3.support import ( + PACKED_EXPERT_DEFAULT_GROUP_SIZE, + PACKED_EXPERT_DEFAULT_NUM_BITS, + get_packed_expert_quantization_args, + has_packed_expert_weights, +) +from xorl.models.transformers.qwen3_5_shared import qwen3_5_apply_rotary_pos_emb +from xorl.ops.fused_silu_and_mul import fused_silu_and_mul +from xorl.utils import logging + + +logger = logging.get_logger(__name__) + + +class DeepseekV3MLP(nn.Module): + def __init__(self, config: DeepseekV3Config, intermediate_size: Optional[int] = None): + super().__init__() + self.hidden_size = config.hidden_size + self.intermediate_size = config.intermediate_size if intermediate_size is None else intermediate_size + self.gate_proj = nn.Linear(config.hidden_size, self.intermediate_size, bias=False) + self.up_proj = nn.Linear(config.hidden_size, self.intermediate_size, bias=False) + self.down_proj = nn.Linear(self.intermediate_size, config.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 forward(self, x: torch.Tensor) -> torch.Tensor: + gate = self.gate_proj(x) + up = self.up_proj(x) + if self._use_fused_silu: + hidden_states = fused_silu_and_mul(torch.cat([gate, up], dim=-1)) + else: + hidden_states = self.act_fn(gate) * up + return self.down_proj(hidden_states) + + +class DeepseekV3TopkRouter(nn.Module): + def __init__(self, config: DeepseekV3Config): + super().__init__() + self.config = config + self.hidden_size = config.hidden_size + self.weight = nn.Parameter(torch.empty(config.n_routed_experts, config.hidden_size)) + self.register_buffer("e_score_correction_bias", torch.zeros(config.n_routed_experts)) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = hidden_states.view(-1, self.hidden_size) + if getattr(self.config, "_router_fp32", False): + return F.linear(hidden_states.float(), self.weight.float()) + return F.linear(hidden_states, self.weight) + + +class DeepseekV3Attention(nn.Module): + def __init__(self, config: DeepseekV3Config, layer_idx: int): + super().__init__() + self.config = config + self.layer_idx = layer_idx + self.num_heads = config.num_attention_heads + self.num_key_value_groups = 1 + self.q_lora_rank = config.q_lora_rank + self.qk_nope_head_dim = config.qk_nope_head_dim + self.qk_rope_head_dim = config.qk_rope_head_dim + self.qk_head_dim = config.qk_head_dim + self.kv_lora_rank = config.kv_lora_rank + self.v_head_dim = config.v_head_dim + self.attention_dropout = config.attention_dropout + self.is_causal = True + self.scaling = self.qk_head_dim**-0.5 + rope_params = config.rope_parameters + if rope_params.get("rope_type", "default") != "default": + mscale_all_dim = rope_params.get("mscale_all_dim", 0) + scaling_factor = rope_params.get("factor") + if mscale_all_dim and scaling_factor is not None: + mscale = 0.1 * mscale_all_dim * torch.log(torch.tensor(float(scaling_factor))).item() + 1.0 + self.scaling = self.scaling * mscale * mscale + + if self.q_lora_rank is None: + self.q_proj = nn.Linear(config.hidden_size, self.num_heads * self.qk_head_dim, bias=False) + else: + self.q_a_proj = nn.Linear(config.hidden_size, config.q_lora_rank, bias=config.attention_bias) + self.q_a_layernorm = RMSNorm(config.q_lora_rank, eps=config.rms_norm_eps) + self.q_b_proj = nn.Linear(config.q_lora_rank, self.num_heads * self.qk_head_dim, bias=False) + + self.kv_a_proj_with_mqa = nn.Linear( + config.hidden_size, + self.kv_lora_rank + self.qk_rope_head_dim, + bias=config.attention_bias, + ) + self.kv_a_layernorm = RMSNorm(self.kv_lora_rank, eps=config.rms_norm_eps) + self.kv_b_proj = nn.Linear( + self.kv_lora_rank, + self.num_heads * (self.qk_nope_head_dim + self.v_head_dim), + bias=False, + ) + self.o_proj = nn.Linear(self.num_heads * self.v_head_dim, config.hidden_size, bias=config.attention_bias) + self._pad_value_for_flash = ( + is_flash_attention(config._attn_implementation) and self.qk_head_dim != self.v_head_dim + ) + + def _project_qkv( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + batch_size, seq_length = hidden_states.shape[:-1] + query_shape = (batch_size, seq_length, self.num_heads, self.qk_head_dim) + kv_shape = (batch_size, seq_length, self.num_heads, self.qk_nope_head_dim + self.v_head_dim) + + if self.q_lora_rank is None: + q_states = self.q_proj(hidden_states) + else: + q_states = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states))) + q_states = q_states.view(query_shape) + q_pass, q_rot = torch.split(q_states, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) + + compressed_kv = self.kv_a_proj_with_mqa(hidden_states) + k_pass, k_rot = torch.split(compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1) + k_pass = self.kv_b_proj(self.kv_a_layernorm(k_pass)).view(kv_shape) + k_pass, value_states = torch.split(k_pass, [self.qk_nope_head_dim, self.v_head_dim], dim=-1) + k_rot = k_rot.view(batch_size, seq_length, 1, self.qk_rope_head_dim) + + cos, sin = position_embeddings + q_rot, k_rot = qwen3_5_apply_rotary_pos_emb( + q_rot, + k_rot, + cos, + sin, + interleaved=getattr(self.config, "rope_interleave", True), + ) + k_rot = k_rot.expand(*k_pass.shape[:-1], -1) + + query_states = torch.cat((q_pass, q_rot), dim=-1) + key_states = torch.cat((k_pass, k_rot), dim=-1) + + if getattr(self.config, "_attention_cast_bf16", False): + query_states = query_states.to(torch.bfloat16) + key_states = key_states.to(torch.bfloat16) + + if self._pad_value_for_flash: + value_states = F.pad(value_states, [0, self.qk_head_dim - self.v_head_dim]) + + return query_states, key_states, value_states + + def _project_output(self, attn_output: torch.Tensor) -> torch.Tensor: + if self._pad_value_for_flash: + attn_output = attn_output[..., : self.v_head_dim] + attn_output = attn_output.reshape(*attn_output.shape[:-2], -1).contiguous() + return self.o_proj(attn_output) + + def _get_attention_fn(self): + return ATTENTION_FUNCTIONS.get(self.config._attn_implementation, eager_attention_forward) + + def _attention_kwargs(self): + return { + "dropout": 0.0 if not self.training else self.attention_dropout, + "scaling": self.scaling, + "sliding_window": None, + } + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + attention_mask: torch.Tensor | None, + **kwargs, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + 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.project_output(self, attn_output) + return attn_output, None + + +class DeepseekV3MoEBlock(MoEBlock): + def __init__(self, config: DeepseekV3Config): + super().__init__( + hidden_size=config.hidden_size, + num_experts=config.n_routed_experts, + top_k=config.num_experts_per_tok, + intermediate_size=config.moe_intermediate_size, + hidden_act=config.hidden_act, + norm_topk_prob=config.norm_topk_prob, + moe_implementation=getattr(config, "_moe_implementation", "eager"), + train_router=getattr(config, "train_router", False), + record_routing_weights=getattr(config, "record_routing_weights", True), + ) + self.config = config + self.gate = DeepseekV3TopkRouter(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.n_group = config.n_group + self.topk_group = config.topk_group + self.norm_topk_prob = config.norm_topk_prob + self.routed_scaling_factor = config.routed_scaling_factor + self.shared_experts = DeepseekV3MLP( + config, + intermediate_size=config.moe_intermediate_size * config.n_shared_experts, + ) + + def _route_tokens_to_experts( + self, + router_logits: torch.Tensor, + input_dtype: torch.dtype, + ) -> tuple[torch.Tensor, torch.Tensor]: + router_scores = router_logits.sigmoid() + choice_scores = router_scores + self.gate.e_score_correction_bias.float() + experts_per_group = self.num_experts // self.n_group + group_topk = min(2, experts_per_group) + group_scores = choice_scores.view(-1, self.n_group, experts_per_group).topk(group_topk, dim=-1)[0].sum(dim=-1) + group_idx = torch.topk(group_scores, k=min(self.topk_group, self.n_group), dim=-1, sorted=False)[1] + group_mask = torch.zeros_like(group_scores) + group_mask.scatter_(1, group_idx, 1) + score_mask = group_mask.unsqueeze(-1).expand(-1, self.n_group, experts_per_group).reshape(-1, self.num_experts) + scores_for_choice = choice_scores.masked_fill(~score_mask.bool(), 0.0) + selected_experts = torch.topk(scores_for_choice, k=self.top_k, dim=-1, sorted=False)[1] + routing_weights = router_scores.gather(1, selected_experts) + if self.norm_topk_prob: + routing_weights = routing_weights / (routing_weights.sum(dim=-1, keepdim=True) + 1e-20) + routing_weights = routing_weights * self.routed_scaling_factor + return routing_weights.to(input_dtype), selected_experts + + def _regather_routing( + self, + router_logits: torch.Tensor, + cached_experts: torch.Tensor, + input_dtype: torch.dtype, + ) -> tuple[torch.Tensor, torch.Tensor]: + routing_weights = torch.gather(router_logits.sigmoid(), 1, cached_experts) + if self.norm_topk_prob: + routing_weights = routing_weights / (routing_weights.sum(dim=-1, keepdim=True) + 1e-20) + routing_weights = routing_weights * self.routed_scaling_factor + return cached_experts, routing_weights.to(input_dtype) + + def forward(self, hidden_states: torch.Tensor): + residuals = hidden_states + batch_size, sequence_length, hidden_dim = hidden_states.shape + flat_hidden_states = hidden_states.view(-1, hidden_dim) + router_logits = self.gate(hidden_states) + + stage = get_replay_stage() + replay = self._routing_replay + + if stage is not None and replay is not None: + if stage == "record": + with torch.no_grad(): + _, selected_experts = self._route_tokens_to_experts(router_logits, flat_hidden_states.dtype) + replay.record(selected_experts) + elif stage == "replay_forward": + selected_experts = replay.pop_forward() + elif stage == "replay_backward": + selected_experts = replay.pop_backward() + else: + raise RuntimeError(f"Unsupported routing replay stage: {stage}") + + selected_experts, routing_weights = self._regather_routing( + router_logits, + selected_experts, + flat_hidden_states.dtype, + ) + + if self.record_routing_weights: + if stage == "record": + replay.record_weights(routing_weights) + elif stage == "replay_backward": + cached_weights = replay.pop_backward_weights() + if cached_weights is not None: + routing_weights = cached_weights.to(flat_hidden_states.dtype) + elif stage == "replay_forward": + cached_weights = replay.pop_forward_weights() + if cached_weights is not None: + routing_weights = cached_weights.to(flat_hidden_states.dtype) + else: + routing_weights, selected_experts = self._route_tokens_to_experts(router_logits, flat_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'." + ) + if not self.train_router: + routing_weights = routing_weights.detach() + + if self.moe_implementation == "eager": + expert_output = self._eager_forward(flat_hidden_states, routing_weights, selected_experts) + else: + expert_output = self.experts(flat_hidden_states, routing_weights, selected_experts) + + expert_output = expert_output.view(batch_size, sequence_length, hidden_dim) + shared_output = self.shared_experts(residuals) + sync_pending_combine() + return expert_output + shared_output, router_logits + + +class DeepseekV3DecoderLayer(nn.Module): + def __init__(self, config: DeepseekV3Config, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + self.self_attn = DeepseekV3Attention(config, layer_idx) + if layer_idx >= config.first_k_dense_replace: + self.mlp = DeepseekV3MoEBlock(config) + else: + self.mlp = DeepseekV3MLP(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, + output_router_logits: Optional[bool] = False, + position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + **kwargs, + ) -> 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 + ) + + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + + if _selective and "self_attn" in self._recompute_modules: + 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 = self.post_attention_layernorm( + hidden_states, + residual=residual, + prenorm=True, + ) + + 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 + + hidden_states = residual + hidden_states + outputs = (hidden_states,) + + if output_attentions: + outputs += (self_attn_weights,) + if output_router_logits: + outputs += (router_logits,) + return outputs + + +class DeepseekV3PreTrainedModel(XorlPreTrainedModel): + config_class = DeepseekV3Config + base_model_prefix = "model" + _no_split_modules = ["DeepseekV3DecoderLayer"] + + 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, DeepseekV3TopkRouter): + module.weight.data.normal_(mean=0.0, std=std) + module.e_score_correction_bias.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_parallel_plan(self): + return parallelize.get_ep_plan() + + def get_checkpoint_handler(self, **kwargs): + checkpoint_keys = kwargs.get("checkpoint_keys", set()) or set() + packed_expert_num_bits = PACKED_EXPERT_DEFAULT_NUM_BITS + packed_expert_group_size = PACKED_EXPERT_DEFAULT_GROUP_SIZE + if checkpoint_keys and has_packed_expert_weights(checkpoint_keys): + packed_quant_args = get_packed_expert_quantization_args(kwargs.get("weights_path", None)) + if packed_quant_args is not None: + packed_expert_num_bits, packed_expert_group_size = packed_quant_args + + ep_rank = kwargs.get("ep_rank", 0) + ep_size = kwargs.get("ep_size", 1) + if kwargs.get("is_broadcast", False): + ep_rank, ep_size = 0, 1 + + return DeepseekV3CheckpointHandler( + num_experts=self.config.n_routed_experts, + ep_rank=ep_rank, + ep_size=ep_size, + packed_expert_num_bits=packed_expert_num_bits, + packed_expert_group_size=packed_expert_group_size, + device=kwargs.get("device"), + dtype=kwargs.get("dtype"), + ) + + +class DeepseekV3Model(DeepseekV3PreTrainedModel): + def __init__(self, config: DeepseekV3Config): + 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( + [DeepseekV3DecoderLayer(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: 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, + output_router_logits: Optional[bool] = None, + **kwargs, + ) -> 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_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 = 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, + sliding_window=None, + is_training=self.training, + output_attentions=output_attentions, + ) + + position_embeddings = self.rotary_emb(hidden_states, position_ids) + ps = get_parallel_state() + position_embeddings = get_cp_strategy().prepare_position_embeddings( + position_embeddings, + dim=1, + sp_group=ps.sp_group, + num_kv_heads=self.config.num_attention_heads, + ) + + all_self_attns = () if output_attentions else None + all_router_logits = () if output_router_logits else None + + for decoder_layer in self.layers: + if decoder_layer is None: + continue + _use_outer_checkpoint = ( + self.gradient_checkpointing and self.training and getattr(self, "_recompute_modules", None) is None + ) + + if _use_outer_checkpoint: + layer_outputs = self._gradient_checkpointing_func( + decoder_layer.__call__, + hidden_states, + causal_mask, + position_ids, + output_attentions, + output_router_logits, + position_embeddings, + **kwargs, + ) + else: + 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, + ) + + hidden_states = layer_outputs[0] + + if output_attentions: + all_self_attns += (layer_outputs[1],) + if output_router_logits and layer_outputs[-1] is not None: + all_router_logits += (layer_outputs[-1],) + + hidden_states = self.norm(hidden_states) if self.norm is not None else hidden_states + if output_hidden_states: + _ = output_hidden_states + + return MoeModelOutput( + last_hidden_state=hidden_states, + attentions=all_self_attns, + router_logits=all_router_logits, + ) + + +class DeepseekV3ForCausalLM(DeepseekV3PreTrainedModel): + _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: DeepseekV3Config): + super().__init__(config) + self.model = DeepseekV3Model(config) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.router_aux_loss_coef = config.router_aux_loss_coef + self.num_experts = config.n_routed_experts + self.num_experts_per_tok = config.num_experts_per_tok + self.post_init() + + 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: 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, + output_router_logits: Optional[bool] = None, + **kwargs, + ) -> MoeCausalLMOutput: + if output_router_logits is None: + output_router_logits = self.config.output_router_logits or self.config.router_aux_loss_coef > 0 + outputs = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + output_router_logits=output_router_logits, + **kwargs, + ) + return MoeCausalLMOutput( + last_hidden_state=outputs.last_hidden_state, + router_logits=outputs.router_logits, + ) + + +ModelClass = DeepseekV3ForCausalLM + + +__all__ = [ + "DeepseekV3ForCausalLM", + "DeepseekV3Model", + "DeepseekV3PreTrainedModel", +] diff --git a/src/xorl/models/transformers/deepseek_v3/parallelize.py b/src/xorl/models/transformers/deepseek_v3/parallelize.py new file mode 100644 index 00000000..4e16b739 --- /dev/null +++ b/src/xorl/models/transformers/deepseek_v3/parallelize.py @@ -0,0 +1,22 @@ +"""Parallelization plan for DeepseekV3 / Kimi-K2.5.""" + +from torch.distributed._tensor import Shard + +from ....distributed.parallel_plan import ParallelPlan + + +MODEL_TP_PLAN = {"lm_head": "colwise_rep"} + + +def get_ep_plan(): + ep_plan = { + "model.layers.*.mlp.experts.gate_up_proj": Shard(0), + "model.layers.*.mlp.experts.down_proj": Shard(0), + "model.layers.*.mlp.experts.gate_proj_lora_A": Shard(0), + "model.layers.*.mlp.experts.gate_proj_lora_B": Shard(0), + "model.layers.*.mlp.experts.up_proj_lora_A": Shard(0), + "model.layers.*.mlp.experts.up_proj_lora_B": Shard(0), + "model.layers.*.mlp.experts.down_proj_lora_A": Shard(0), + "model.layers.*.mlp.experts.down_proj_lora_B": Shard(0), + } + return ParallelPlan(ep_plan=ep_plan) diff --git a/src/xorl/models/transformers/deepseek_v3/support.py b/src/xorl/models/transformers/deepseek_v3/support.py new file mode 100644 index 00000000..ac197d35 --- /dev/null +++ b/src/xorl/models/transformers/deepseek_v3/support.py @@ -0,0 +1,138 @@ +"""DeepseekV3 / Kimi-K2.5 support helpers.""" + +import json +import os +from typing import Iterable, Optional, Tuple + +from transformers.utils import cached_file + +from xorl.distributed.parallel_state import get_parallel_state + + +PACKED_EXPERT_DEFAULT_NUM_BITS = 4 +PACKED_EXPERT_DEFAULT_GROUP_SIZE = 32 + + +DEEPSEEK_V3_LORA_TARGET_MODULES = [ + "q_a_proj", + "q_b_proj", + "kv_a_proj_with_mqa", + "kv_b_proj", + "o_proj", + "gate_proj", + "up_proj", + "down_proj", +] + + +def is_deepseek_v3_config(config) -> bool: + return getattr(config, "model_type", None) == "deepseek_v3" + + +def validate_deepseek_v3_router_settings(config, *, train_router: bool) -> None: + if not is_deepseek_v3_config(config): + return + if train_router: + raise ValueError("DeepseekV3/Kimi-K2.5 does not support train_router=True in xorl yet.") + + +def validate_deepseek_v3_training_mode( + config, + *, + enable_qlora: bool, + freeze_router: bool, + merge_qkv: bool, +) -> None: + if not is_deepseek_v3_config(config): + return + if enable_qlora: + raise ValueError("DeepseekV3/Kimi-K2.5 does not support enable_qlora=True yet.") + if not freeze_router: + raise ValueError("DeepseekV3/Kimi-K2.5 requires freeze_router=True.") + if not merge_qkv: + raise ValueError("DeepseekV3/Kimi-K2.5 does not support merge_qkv=False yet.") + + +def validate_deepseek_v3_tensor_parallelism(config) -> None: + if not is_deepseek_v3_config(config): + return + if get_parallel_state().tp_enabled: + raise ValueError("DeepseekV3/Kimi-K2.5 tensor parallelism is not supported yet.") + + +def freeze_deepseek_v3_router_parameters(model) -> int: + count = 0 + for name, param in model.named_parameters(): + if ".gate.weight" in name: + param.requires_grad = False + count += 1 + return count + + +def deepseek_v3_default_lora_targets(*, train_attn: bool, train_mlp: bool, train_unembed: bool) -> list[str]: + targets: list[str] = [] + if train_attn: + targets.extend(DEEPSEEK_V3_LORA_TARGET_MODULES[:5]) + if train_mlp: + targets.extend(DEEPSEEK_V3_LORA_TARGET_MODULES[5:]) + if train_unembed: + targets.append("lm_head") + return targets + + +def has_packed_expert_weights(checkpoint_keys: Iterable[str]) -> bool: + return any(".mlp.experts." in key and key.endswith(".weight_packed") for key in checkpoint_keys) + + +def _resolve_weights_path(weights_path: Optional[str]) -> Optional[str]: + if not weights_path: + return None + if os.path.isdir(weights_path): + return weights_path + try: + config_path = cached_file(weights_path, "config.json", _raise_exceptions_for_missing_entries=False) + except Exception: + return None + if config_path and os.path.isfile(config_path): + return os.path.dirname(config_path) + return None + + +def get_packed_expert_quantization_args(weights_path: Optional[str]) -> Optional[Tuple[int, int]]: + resolved_weights_path = _resolve_weights_path(weights_path) + if resolved_weights_path is None: + return None + + config_path = os.path.join(resolved_weights_path, "config.json") + if not os.path.isfile(config_path): + return None + + try: + with open(config_path) as f: + config_dict = json.load(f) + except (json.JSONDecodeError, OSError): + return None + + quantization_config = config_dict.get("quantization_config") + if quantization_config is None and isinstance(config_dict.get("text_config"), dict): + quantization_config = config_dict["text_config"].get("quantization_config") + if not isinstance(quantization_config, dict): + return None + if quantization_config.get("quant_method") != "compressed-tensors": + return None + if quantization_config.get("format") != "pack-quantized": + return None + + config_groups = quantization_config.get("config_groups", {}) + for group_config in config_groups.values(): + if not isinstance(group_config, dict): + continue + weights_config = group_config.get("weights") + if not isinstance(weights_config, dict): + continue + return ( + int(weights_config.get("num_bits", PACKED_EXPERT_DEFAULT_NUM_BITS)), + int(weights_config.get("group_size", PACKED_EXPERT_DEFAULT_GROUP_SIZE)), + ) + + return PACKED_EXPERT_DEFAULT_NUM_BITS, PACKED_EXPERT_DEFAULT_GROUP_SIZE diff --git a/src/xorl/models/transformers/deepseek_v3/tokenization_kimi.py b/src/xorl/models/transformers/deepseek_v3/tokenization_kimi.py new file mode 100644 index 00000000..9d102557 --- /dev/null +++ b/src/xorl/models/transformers/deepseek_v3/tokenization_kimi.py @@ -0,0 +1,248 @@ +import os +from collections import OrderedDict +from logging import getLogger +from pathlib import Path +from shutil import copyfile +from typing import Dict, Iterator, List, Optional, Tuple, Union, cast + +import tiktoken +from tiktoken.load import load_tiktoken_bpe +from tokenizers import AddedToken +from transformers.tokenization_utils import PreTrainedTokenizer + + +try: + from transformers.convert_slow_tokenizer import bytes_to_unicode +except ImportError: # pragma: no cover - compatibility with older Transformers + from transformers.models.gpt2.tokenization_gpt2 import bytes_to_unicode + + +logger = getLogger(__name__) +VOCAB_FILES_NAMES = {"vocab_file": "tiktoken.model"} + + +def _normalize_added_tokens_decoder(added_tokens_decoder: Optional[dict]) -> Optional[dict[int, AddedToken]]: + if added_tokens_decoder is None: + return None + + normalized = {} + for token_id, token_config in added_tokens_decoder.items(): + token_id = int(token_id) + if isinstance(token_config, AddedToken): + normalized[token_id] = token_config + elif isinstance(token_config, dict): + normalized[token_id] = AddedToken( + token_config["content"], + single_word=token_config.get("single_word", False), + lstrip=token_config.get("lstrip", False), + rstrip=token_config.get("rstrip", False), + normalized=token_config.get("normalized", False), + special=token_config.get("special", False), + ) + else: + normalized[token_id] = AddedToken(str(token_config), special=True) + return normalized + + +class TikTokenTokenizer(PreTrainedTokenizer): + """Local Kimi TikToken tokenizer. + + Moonshot Kimi snapshots publish this tokenizer as remote HuggingFace code. + Xorl vendors the small tokenizer implementation locally so loading pinned + local snapshots does not require enabling HuggingFace remote code. + """ + + vocab_files_names = VOCAB_FILES_NAMES + model_input_names = ["input_ids", "attention_mask"] + num_reserved_special_tokens = 256 + + special_tokens: Dict[str, int] + + pat_str = "|".join( + [ + r"""[\p{Han}]+""", + r"""[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]*[\p{Ll}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]+(?i:'s|'t|'re|'ve|'m|'ll|'d)?""", + r"""[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]+[\p{Ll}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]*(?i:'s|'t|'re|'ve|'m|'ll|'d)?""", + r"""\p{N}{1,3}""", + r""" ?[^\s\p{L}\p{N}]+[\r\n]*""", + r"""\s*[\r\n]+""", + r"""\s+(?!\S)""", + r"""\s+""", + ] + ) + + def __init__( + self, + vocab_file: str, + bos_token: Union[str, AddedToken] = "[BOS]", + eos_token: Union[str, AddedToken] = "[EOS]", + unk_token: Optional[Union[str, AddedToken]] = None, + pad_token: Optional[Union[str, AddedToken]] = None, + additional_special_tokens: Optional[List[str]] = None, + added_tokens_decoder: Optional[dict] = None, + **kwargs, + ): + if not os.path.isfile(vocab_file): + raise FileNotFoundError(vocab_file) + + if additional_special_tokens is None: + additional_special_tokens = [ + "<|im_end|>", + "<|im_user|>", + "<|im_assistant|>", + "<|start_header_id|>", + "<|end_header_id|>", + "[EOT]", + "<|im_system|>", + "<|im_middle|>", + ] + + added_tokens_decoder = _normalize_added_tokens_decoder(added_tokens_decoder) + special_tokens_mapping = {} + if added_tokens_decoder: + special_tokens_mapping = {token_id: token.content for token_id, token in added_tokens_decoder.items()} + + self.vocab_file = vocab_file + mergeable_ranks = load_tiktoken_bpe(vocab_file) + num_base_tokens = len(mergeable_ranks) + self.special_tokens = { + special_tokens_mapping.get(i, f"<|reserved_token_{i}|>"): i + for i in range(num_base_tokens, num_base_tokens + self.num_reserved_special_tokens) + } + + self.model = tiktoken.Encoding( + name=Path(vocab_file).name, + pat_str=self.pat_str, + mergeable_ranks=mergeable_ranks, + special_tokens=self.special_tokens, + ) + logger.info(f"Loaded local Kimi tiktoken model from {vocab_file}") + + self.n_words: int = self.model.n_vocab + self.bos_id: int = self.special_tokens[str(bos_token)] + self.eos_id: int = self.special_tokens[str(eos_token)] + self.pad_id: Optional[int] = self.special_tokens[str(pad_token)] if pad_token is not None else None + self.unk_id: Optional[int] = self.special_tokens[str(unk_token)] if unk_token is not None else None + + self.byte_encoder = bytes_to_unicode() + self.byte_decoder = {v: k for k, v in self.byte_encoder.items()} + self.decoder = {} + special_token_ids = set(self.special_tokens.values()) + for i in range(self.n_words): + if i in special_token_ids: + continue + decoding = "".join( + self.byte_encoder[ord(char)] for char in self.model.decode_single_token_bytes(i).decode("latin-1") + ) + self.decoder[i] = decoding + self.decoder.update({token_id: token for token, token_id in self.special_tokens.items()}) + + self.encoder = {token: token_id for token_id, token in self.decoder.items()} + self._token_config_cache = OrderedDict() + self._cache_max_size = 128 + + super().__init__( + bos_token=bos_token, + eos_token=eos_token, + unk_token=unk_token, + pad_token=pad_token, + additional_special_tokens=additional_special_tokens, + added_tokens_decoder=added_tokens_decoder, + **kwargs, + ) + self.all_special_ids_set = set(self.all_special_ids) + + def encode(self, text: str, allow_special_tokens: bool = True, **kwargs) -> List[int]: + if len(kwargs) > 0: + logger.warning(f"Calling super().encode with {kwargs}") + return super().encode(text, **kwargs) + + if not isinstance(text, str): + raise TypeError(f"text must be str, got {type(text).__name__}") + + tiktoken_max_encode_chars = 400_000 + max_no_whitespaces_chars = 25_000 + substrs = [] + for processed_text in self.pre_tokenizer_process(text): + substrs.extend( + substr + for i in range(0, len(processed_text), tiktoken_max_encode_chars) + for substr in self._split_whitespaces_or_nonwhitespaces( + processed_text[i : i + tiktoken_max_encode_chars], + max_no_whitespaces_chars, + ) + ) + + token_ids: List[int] = [] + for substr in substrs: + if allow_special_tokens: + token_ids.extend(self.model.encode(substr, allowed_special="all")) + else: + token_ids.extend(self.model.encode(substr, disallowed_special=())) + return token_ids + + def decode(self, token_ids: Union[int, List[int]], **kwargs) -> str: + if len(kwargs) > 0: + return super().decode(token_ids, **kwargs) + if isinstance(token_ids, int): + token_ids = [token_ids] + return self.model.decode(cast(List[int], token_ids)) + + @staticmethod + def _split_whitespaces_or_nonwhitespaces(s: str, max_consecutive_slice_len: int) -> Iterator[str]: + current_slice_len = 0 + current_slice_is_space = s[0].isspace() if len(s) > 0 else False + slice_start = 0 + + for i, char in enumerate(s): + is_now_space = char.isspace() + if current_slice_is_space ^ is_now_space: + current_slice_len = 1 + current_slice_is_space = is_now_space + else: + current_slice_len += 1 + if current_slice_len > max_consecutive_slice_len: + yield s[slice_start:i] + slice_start = i + current_slice_len = 1 + yield s[slice_start:] + + def pre_tokenizer_process(self, text: str) -> List[str]: + return [text] + + @property + def vocab_size(self) -> int: + return self.n_words + + def get_vocab(self) -> Dict[str, int]: + return dict(self.encoder) + + def _tokenize(self, text: str, **kwargs) -> List[str]: + return [self.decoder[t] for t in self.encode(text)] + + def _convert_token_to_id(self, token: str) -> Optional[int]: + return self.encoder.get(token, self.unk_id) + + def _convert_id_to_token(self, index: int) -> Optional[str]: + return self.decoder.get(index) + + @staticmethod + def clean_up_tokenization(out_string: str) -> str: + return out_string + + def convert_tokens_to_string(self, tokens: List[str]) -> str: + text = "".join(tokens) + return bytearray([self.byte_decoder[c] for c in text]).decode("utf-8", "replace") + + def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]: + if not os.path.isdir(save_directory): + raise ValueError(f"vocabulary path ({save_directory}) should be a directory") + out_vocab_file = os.path.join( + save_directory, + (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"], + ) + + if os.path.abspath(self.vocab_file) != os.path.abspath(out_vocab_file) and os.path.isfile(self.vocab_file): + copyfile(self.vocab_file, out_vocab_file) + + return (out_vocab_file,) diff --git a/src/xorl/models/transformers/glm4_moe/__init__.py b/src/xorl/models/transformers/glm4_moe/__init__.py new file mode 100644 index 00000000..ee99735c --- /dev/null +++ b/src/xorl/models/transformers/glm4_moe/__init__.py @@ -0,0 +1,38 @@ +"""GLM-4 MoE model.""" + +from ...layers.moe import ( + MOE_EXPERT_BACKENDS, + MoEBlock, + MoEExperts, + MoEExpertsLoRA, + MoELoRAConfig, + TopKRouter, +) +from .configuration_glm4_moe import Glm4MoeConfig +from .modeling_glm4_moe import ( + Glm4MoeAttention, + Glm4MoeForCausalLM, + Glm4MoeGate, + Glm4MoeMLP, + Glm4MoeModel, + Glm4MoePreTrainedModel, + Glm4MoeSparseMoeBlock, +) + + +__all__ = [ + "Glm4MoeConfig", + "Glm4MoeForCausalLM", + "Glm4MoeModel", + "Glm4MoePreTrainedModel", + "Glm4MoeSparseMoeBlock", + "Glm4MoeGate", + "Glm4MoeMLP", + "Glm4MoeAttention", + "MoEBlock", + "MoEExperts", + "MoEExpertsLoRA", + "MoELoRAConfig", + "TopKRouter", + "MOE_EXPERT_BACKENDS", +] diff --git a/src/xorl/models/transformers/glm4_moe/checkpoint_handler.py b/src/xorl/models/transformers/glm4_moe/checkpoint_handler.py new file mode 100644 index 00000000..c31af015 --- /dev/null +++ b/src/xorl/models/transformers/glm4_moe/checkpoint_handler.py @@ -0,0 +1,345 @@ +"""Checkpoint handler for GLM-4 MoE models. + +Reuses the same buffer infrastructure as Qwen3 MoE: +- ExpertWeightBuffer: stacks per-expert HF weights into [num_experts, ...] tensors +- QKVMergeBuffer: merges q_proj + k_proj + v_proj -> qkv_proj +- GateUpMergeBuffer: merges gate_proj + up_proj -> gate_up_proj (dense + shared experts) + +GLM-specific: ``gate.e_score_correction_bias`` passes through directly (checkpoint +path matches model path). +""" + +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, + EXPERT_QUANT_AUX_PATTERN, + FP8_AUX_SUFFIX_PATTERN, + OPROJ_WEIGHT_PATTERN, + QKV_PROJ_PATTERN, + QUANT_AUX_SUFFIX_PATTERN, + ExpertWeightBuffer, + GateUpMergeBuffer, + QKVMergeBuffer, + QLoRAExpertBuffer, + QLoRAWeightBuffer, + parse_expert_full_key, + parse_expert_key, +) + + +class Glm4MoeCheckpointHandler(CheckpointHandler): + """Checkpoint handler for GLM-4 MoE models. + + Load transforms: + 1. Per-expert weights -> fused [num_experts, ...] stacked tensors + 2. Dense layer / shared expert gate_proj + up_proj -> gate_up_proj + 3. q_proj + k_proj + v_proj -> qkv_proj + + Save transforms: + 1. gate_up_proj -> gate_proj + up_proj + 2. qkv_proj -> q_proj + k_proj + v_proj + 3. Expert tensors pass through as-is (fused format) + """ + + def __init__( + self, + num_experts: int, + num_attention_heads: int, + num_key_value_heads: int, + head_dim: int, + ep_rank: int = 0, + ep_size: int = 1, + checkpoint_has_per_expert: bool = True, + skip_qkv_merge: bool = False, + skip_gate_up_merge: bool = False, + is_prequantized: bool = False, + exclude_modules: Optional[Set[str]] = None, + device: Optional["torch.device"] = None, + model: Optional[nn.Module] = None, + num_hidden_layers: Optional[int] = None, + ): + self._expert_buffer: Optional[ExpertWeightBuffer] = None + if checkpoint_has_per_expert and not is_prequantized: + self._expert_buffer = ExpertWeightBuffer( + num_experts, + ep_rank=ep_rank, + ep_size=ep_size, + device=device, + ) + self._qkv_buffer: Optional[QKVMergeBuffer] = None + if not skip_qkv_merge: + self._qkv_buffer = QKVMergeBuffer() + self._gate_up_buffer: Optional[GateUpMergeBuffer] = None + if not skip_gate_up_merge: + self._gate_up_buffer = GateUpMergeBuffer() + 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() + # MTP (multi-token prediction) layer remapping: GLM-4.7 stores embedding, + # output norm, and LM head under model.layers.{num_hidden_layers} in the + # checkpoint, but the model expects them at top-level positions. + self._mtp_layer_prefix = f"model.layers.{num_hidden_layers}." + self._mtp_remap = { + f"model.layers.{num_hidden_layers}.embed_tokens.weight": "model.embed_tokens.weight", + f"model.layers.{num_hidden_layers}.shared_head.norm.weight": "model.norm.weight", + f"model.layers.{num_hidden_layers}.shared_head.head.weight": "lm_head.weight", + } + self._qlora_buffer: Optional[QLoRAWeightBuffer] = None + if is_prequantized and model is not None: + self._qlora_buffer = QLoRAWeightBuffer(model) + 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, + ) + + def get_skip_key_fn(self) -> Optional[Callable[[str], bool]]: + 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 + ) + + if not has_ep_filter and not has_expert_ep_filter and not self._is_prequantized: + 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 + 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 + if has_qlora_expert_buffer: + qe_start = self._qlora_expert_buffer.expert_start + qe_end = self._qlora_expert_buffer.expert_end + + def _should_skip(key: str) -> bool: + if is_prequantized: + 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 has_qlora_expert_buffer: + parsed = parse_expert_full_key(key) + if parsed is not None: + _, expert_idx, _, suffix = parsed + if suffix == "input_scale": + return True + return expert_idx < qe_start or expert_idx >= qe_end + else: + if parse_expert_key(key) is not None: + return True + if EXPERT_QUANT_AUX_PATTERN.match(key) is not None: + return True + + if not has_qlora_buffer: + 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 + + if has_ep_filter: + parsed = parse_expert_key(key) + if parsed is not None: + _, expert_idx, _ = parsed + return expert_idx < ep_start or expert_idx >= ep_end + return False + + return _should_skip + + def _maybe_finalize_per_expert_merge( + self, + layer_idx: int, + proj: str, + ) -> List[Tuple[str, torch.Tensor]]: + """Finalize expert weight merging: fuse gate+up into gate_up_proj.""" + 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 _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]]: + if key.endswith(".input_scale"): + return [] + + # QLoRA expert buffer + if self._qlora_expert_buffer is not None and not self._is_excluded_module(key): + result = self._qlora_expert_buffer.try_consume(key, tensor) + if result is not None: + return result + + # QLoRA buffer + 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 + + if self._is_prequantized and not self._is_excluded_module(key): + if QUANT_AUX_SUFFIX_PATTERN.search(key): + return [] + if FP8_AUX_SUFFIX_PATTERN.search(key): + return [] + if parse_expert_key(key) is not None: + return [] + if EXPERT_QUANT_AUX_PATTERN.match(key) is not None: + return [] + if key.endswith(".weight"): + if OPROJ_WEIGHT_PATTERN.match(key) or DENSE_DOWN_PROJ_PATTERN.match(key): + return [] + + # 1. MTP layer remapping: remap shared output components, skip MTP-only weights. + # Must be checked first to prevent MTP layer expert/QKV/gate_up keys from + # entering the merge buffers (which would emit fused keys for a nonexistent layer). + if key.startswith(self._mtp_layer_prefix): + if key in self._mtp_remap: + return [(self._mtp_remap[key], tensor)] + # Skip MTP-only components (eh_proj, enorm, hnorm, and MTP decoder layer) + return [] + + # 2. Expert merge + 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) + result = self._maybe_finalize_per_expert_merge(layer_idx, proj) + return result if result else [] + + # 3. QKV merge + if self._qkv_buffer is not None: + if self._is_prequantized and key.endswith(".weight"): + if self._qkv_buffer.is_qkv_key(key): + return [] + + 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 [] + + # 4. Gate/up merge + if self._gate_up_buffer is not None: + if self._is_prequantized and key.endswith(".weight"): + if self._gate_up_buffer.is_gate_up_key(key): + return [] + + 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 [] + + # 5. Passthrough (includes gate.e_score_correction_bias) + return [(key, tensor)] + + def on_skip_weight(self, key: str) -> List[Tuple[str, torch.Tensor]]: + if self._qlora_expert_buffer is not None: + self._qlora_expert_buffer.count_skipped(key) + + 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) + result = self._maybe_finalize_per_expert_merge(layer_idx, proj) + if result: + return result + return [] + + def on_load_complete(self) -> List[Tuple[str, torch.Tensor]]: + 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._gate_up_buffer is not None: + pending_gu = self._gate_up_buffer.get_pending() + if pending_gu: + warnings.warn(f"Incomplete gate/up merge pairs after loading: {pending_gu}") + if self._qkv_buffer is not None: + 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() + if self._qlora_expert_buffer is not None: + pending_exp = self._qlora_expert_buffer.get_pending() + if pending_exp: + warnings.warn( + 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]]: + # Split gate_up_proj -> gate_proj + up_proj (dense / shared experts) + 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/glm4_moe/configuration_glm4_moe.py b/src/xorl/models/transformers/glm4_moe/configuration_glm4_moe.py new file mode 100644 index 00000000..ff661a28 --- /dev/null +++ b/src/xorl/models/transformers/glm4_moe/configuration_glm4_moe.py @@ -0,0 +1,174 @@ +# Copyright 2025 The ZhipuAI Inc. team and the HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""GLM-4 MoE model configuration""" + +from transformers.configuration_utils import PretrainedConfig + +from xorl.models.layers import rope_config_validation + +from ....utils import logging +from .parallelize import TP_PLAN + + +logger = logging.get_logger(__name__) + + +class Glm4MoeConfig(PretrainedConfig): + r""" + Configuration class for the GLM-4 MoE model (GLM-4.5 / 4.6 / 4.7 variants). + + Defaults correspond to + `THUDM/GLM-4-0414-A10B-Base `_. + + Args: + vocab_size: Vocabulary size. + hidden_size: Dimension of the hidden representations. + intermediate_size: MLP dimension for dense layers (first K layers). + num_hidden_layers: Number of transformer layers. + num_attention_heads: Number of attention heads. + num_key_value_heads: Number of key/value heads for GQA. + hidden_act: Activation function name. + max_position_embeddings: Maximum sequence length. + initializer_range: Standard deviation for weight initialization. + rms_norm_eps: Epsilon for RMSNorm. + use_cache: Whether to return past key/value states. + tie_word_embeddings: Whether to tie input/output embeddings. + rope_theta: Base period for RoPE. + rope_scaling: RoPE scaling configuration dict. + partial_rotary_factor: Fraction of head dimensions that receive rotary embeddings. + attention_bias: Whether to use bias in QKV projections. + attention_dropout: Dropout rate for attention weights. + moe_intermediate_size: Expert FFN intermediate dimension. + num_experts_per_tok: Number of experts selected per token. + n_shared_experts: Number of shared (dense) experts alongside routed experts. + n_routed_experts: Total number of routed experts. + routed_scaling_factor: Multiplicative factor applied to routed expert outputs. + n_group: Number of expert groups for grouped top-k routing. + topk_group: Number of groups selected per token before picking top-k within groups. + first_k_dense_replace: First K layers use dense MLP instead of MoE. + norm_topk_prob: Whether to renormalize top-k routing weights. + use_qk_norm: Whether to apply per-head RMSNorm to Q and K. + output_router_logits: Whether to return router logits from the model. + router_aux_loss_coef: Coefficient for auxiliary load-balancing loss. + _moe_implementation: MoE backend name (``"triton"``, ``"quack"``, ``"native"``, ``"eager"``). + """ + + model_type = "glm4_moe" + + 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=151552, + hidden_size=4096, + intermediate_size=10944, + num_hidden_layers=46, + num_attention_heads=96, + num_key_value_heads=8, + hidden_act="silu", + max_position_embeddings=131072, + initializer_range=0.02, + rms_norm_eps=1e-5, + use_cache=False, + tie_word_embeddings=False, + rope_theta=10000.0, + rope_scaling=None, + partial_rotary_factor=0.5, + attention_bias=False, + attention_dropout=0.0, + moe_intermediate_size=1408, + num_experts_per_tok=8, + n_shared_experts=1, + n_routed_experts=128, + routed_scaling_factor=1.0, + n_group=1, + topk_group=1, + first_k_dense_replace=1, + norm_topk_prob=True, + use_qk_norm=False, + output_router_logits=False, + router_aux_loss_coef=0.001, + _moe_implementation="triton", + **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.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.partial_rotary_factor = partial_rotary_factor + self.attention_bias = attention_bias + self.attention_dropout = attention_dropout + self.use_qk_norm = use_qk_norm + + 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) + + # MoE arguments + self.moe_intermediate_size = moe_intermediate_size + self.num_experts_per_tok = num_experts_per_tok + self.n_group = n_group + self.topk_group = topk_group + self.n_shared_experts = n_shared_experts + self.n_routed_experts = n_routed_experts + self.routed_scaling_factor = routed_scaling_factor + self.first_k_dense_replace = first_k_dense_replace + self.norm_topk_prob = norm_topk_prob + self.output_router_logits = output_router_logits + self.router_aux_loss_coef = router_aux_loss_coef + self._moe_implementation = _moe_implementation + + super().__init__( + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) + + @property + def rope_parameters(self): + """Return rope parameters exposing partial_rotary_factor for RotaryEmbedding.""" + rope_params = { + "rope_type": "default", + "rope_theta": self.rope_theta, + "partial_rotary_factor": self.partial_rotary_factor, + } + 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") + return rope_params + + @rope_parameters.setter + def rope_parameters(self, value): + """Setter for rope_parameters to satisfy HF transformers 5.0+ configuration.""" + if value is not None and isinstance(value, dict): + if "rope_theta" in value: + self.rope_theta = value["rope_theta"] + self._rope_scaling = value + + +__all__ = ["Glm4MoeConfig"] diff --git a/src/xorl/models/transformers/glm4_moe/modeling_glm4_moe.py b/src/xorl/models/transformers/glm4_moe/modeling_glm4_moe.py new file mode 100644 index 00000000..5446b0c0 --- /dev/null +++ b/src/xorl/models/transformers/glm4_moe/modeling_glm4_moe.py @@ -0,0 +1,721 @@ +# Copyright 2025 The ZhipuAI Inc. team and the HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Optional, Tuple, Unpack + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from xorl.distributed.moe.deepep import sync_pending_combine +from xorl.distributed.parallel_state import get_parallel_state +from xorl.distributed.sequence_parallel.strategy import get_cp_strategy +from xorl.models.base import XorlPreTrainedModel +from xorl.models.checkpoint_handlers.buffers import ( + checkpoint_has_per_expert_weights, + detect_prequantized_block_fp8_checkpoint, + detect_prequantized_checkpoint, + get_prequantized_exclude_modules, +) +from xorl.models.layers import ACT2FN, RotaryEmbedding +from xorl.models.layers.attention import ( + AttentionKwargs, + MultiHeadAttention, + is_flash_attention, + update_causal_mask, +) +from xorl.models.layers.moe import MoEBlock +from xorl.models.layers.moe.routing_replay import get_replay_stage +from xorl.models.layers.normalization import compiled_eager_rms_norm +from xorl.models.layers.rope import apply_rotary_pos_emb +from xorl.models.outputs import MoeCausalLMOutput, MoeModelOutput +from xorl.models.transformers.glm4_moe import parallelize +from xorl.models.transformers.glm4_moe.checkpoint_handler import Glm4MoeCheckpointHandler +from xorl.models.transformers.glm4_moe.configuration_glm4_moe import Glm4MoeConfig +from xorl.ops.fused_silu_and_mul import fused_silu_and_mul +from xorl.utils import logging + + +logger = logging.get_logger(__name__) + + +# --------------------------------------------------------------------------- +# RMSNorm — hardcoded fp32 upcast to match HF Glm4MoeRMSNorm. +# Bypasses the global rmsnorm_mode switch because the bf16 path +# (F.rms_norm) compounds errors across GLM-4 MoE's 92+ norm layers +# and interacts pathologically with the triton/quack MoE kernels — +# observed as 23x logprob spikes vs SGLang/HF on borderline routing. +# +# The forward path delegates to ``compiled_eager_rms_norm`` so the 5 +# pointwise/reduction kernels of the fp32-upcast eager path get fused +# into a single Inductor kernel — important because GLM-4 MoE has 92+ +# norm layers per forward pass. +# --------------------------------------------------------------------------- + + +class Glm4MoeRMSNorm(nn.Module): + def __init__(self, hidden_size: int, eps: float = 1e-6): + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return compiled_eager_rms_norm(hidden_states, self.weight, self.variance_epsilon) + + def extra_repr(self) -> str: + return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}" + + +# --------------------------------------------------------------------------- +# Dense MLP (used for first K layers + shared experts) +# --------------------------------------------------------------------------- + + +class Glm4MoeMLP(nn.Module): + def __init__(self, config, intermediate_size=None): + super().__init__() + self.hidden_size = config.hidden_size + self.intermediate_size = intermediate_size if intermediate_size is not None else 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.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) + + def unfuse_for_tp(self): + 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) + + +# --------------------------------------------------------------------------- +# Router gate with correction bias (sigmoid-based grouped top-k) +# --------------------------------------------------------------------------- + + +class Glm4MoeGate(nn.Module): + """Router gate with e_score_correction_bias for GLM-4 MoE. + + Checkpoint paths: + - ``mlp.gate.weight`` -> ``(n_routed_experts, hidden_size)`` + - ``mlp.gate.e_score_correction_bias`` -> ``(n_routed_experts,)`` + """ + + def __init__(self, hidden_size: int, n_routed_experts: int): + super().__init__() + self.weight = nn.Parameter(torch.empty(n_routed_experts, hidden_size)) + self.e_score_correction_bias = nn.Parameter(torch.zeros(n_routed_experts), requires_grad=False) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return F.linear(hidden_states.float(), self.weight.float()) + + +# --------------------------------------------------------------------------- +# MoE block with sigmoid grouped top-k routing + shared experts +# --------------------------------------------------------------------------- + + +class Glm4MoeSparseMoeBlock(MoEBlock): + """GLM-4 MoE block: sigmoid routing, grouped top-k, correction bias, shared experts. + + Inherits from ``MoEBlock`` to get routing replay and ``moe_act`` gradient + checkpointing support from ``XorlPreTrainedModel``. + + Overrides ``__init__`` and ``forward`` for GLM-specific routing. + """ + + def __init__(self, config: Glm4MoeConfig, moe_implementation: str = "triton"): + super().__init__( + hidden_size=config.hidden_size, + num_experts=config.n_routed_experts, + top_k=config.num_experts_per_tok, + intermediate_size=config.moe_intermediate_size, + hidden_act=config.hidden_act, + norm_topk_prob=config.norm_topk_prob, + moe_implementation=moe_implementation, + ) + self.config = config + self.n_group = config.n_group + self.topk_group = config.topk_group + self.routed_scaling_factor = config.routed_scaling_factor + + # Replace the default nn.Linear gate with Glm4MoeGate (has correction bias) + del self.gate + self.gate = Glm4MoeGate(config.hidden_size, config.n_routed_experts) + + # Shared experts (dense MLP alongside routed experts) + if config.n_shared_experts > 0: + shared_intermediate = config.moe_intermediate_size * config.n_shared_experts + self.shared_experts = Glm4MoeMLP(config, intermediate_size=shared_intermediate) + else: + self.shared_experts = None + + 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) + + def _route_tokens(self, router_logits: torch.Tensor, input_dtype: torch.dtype): + """Sigmoid-based grouped top-k routing. + + 1. Compute sigmoid scores; add correction bias for expert *selection* only + 2. Group experts, select top groups per token + 3. Select top-k experts within selected groups + 4. Gather routing weights from *raw* sigmoid scores (without bias) + 5. Optionally normalize and apply scaling factor + """ + scores = router_logits.sigmoid() + scores_for_choice = scores + self.gate.e_score_correction_bias.unsqueeze(0) + + if self.n_group > 1: + scores_grouped = scores_for_choice.view(scores.shape[0], self.n_group, -1) + group_scores = scores_grouped.topk(2, dim=-1)[0].sum(dim=-1) + group_idx = torch.topk(group_scores, k=self.topk_group, dim=-1, sorted=False)[1] + group_mask = torch.zeros_like(group_scores) + group_mask.scatter_(1, group_idx, 1) + score_mask = group_mask.unsqueeze(-1).expand(-1, -1, scores_grouped.shape[-1]).reshape(scores.shape[0], -1) + scores_for_choice = scores_for_choice.masked_fill(~score_mask.bool(), float("-inf")) + _, selected_experts = torch.topk(scores_for_choice, self.top_k, dim=-1) + routing_weights = scores.gather(1, selected_experts) + + if self.router.norm_topk_prob: + routing_weights = routing_weights / routing_weights.sum(dim=-1, keepdim=True) + routing_weights = routing_weights * self.routed_scaling_factor + routing_weights = routing_weights.to(input_dtype) + return routing_weights, selected_experts + + def _regather_routing(self, router_logits, cached_experts, input_dtype): + """Re-gather routing weights from raw sigmoid scores using cached expert indices. + + The correction bias is only used for expert *selection* (top-k), not for + computing the routing weights themselves. + """ + scores = router_logits.sigmoid() + routing_weights = torch.gather(scores, 1, cached_experts) + if self.router.norm_topk_prob: + routing_weights = routing_weights / routing_weights.sum(dim=-1, keepdim=True) + routing_weights = routing_weights * self.routed_scaling_factor + routing_weights = routing_weights.to(input_dtype) + return cached_experts, routing_weights + + def forward(self, hidden_states: torch.Tensor): + batch_size, sequence_length, hidden_dim = hidden_states.shape + hidden_states_flat = hidden_states.view(-1, hidden_dim) + + router_logits = self.gate(hidden_states_flat) + + stage = get_replay_stage() + replay = self._routing_replay + + if stage is not None and replay is not None: + cached_weights = None + if stage == "record": + with torch.no_grad(): + _, selected_experts = self._route_tokens(router_logits, hidden_states_flat.dtype) + 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: + routing_weights = cached_weights.to(hidden_states_flat.dtype) + else: + selected_experts, routing_weights = self._regather_routing( + router_logits, selected_experts, hidden_states_flat.dtype + ) + else: + routing_weights, selected_experts = self._route_tokens(router_logits, hidden_states_flat.dtype) + + if not self.train_router: + routing_weights = routing_weights.detach() + + # Expert computation + if self.moe_implementation == "eager": + routed_output = self._eager_forward(hidden_states_flat, routing_weights, selected_experts) + else: + routed_output = self.experts(hidden_states_flat, routing_weights, selected_experts) + + # Shared expert computation + if self.shared_experts is not None: + shared_output = self.shared_experts(hidden_states_flat) + final_hidden_states = routed_output + shared_output + else: + final_hidden_states = routed_output + + final_hidden_states = final_hidden_states.reshape(batch_size, sequence_length, hidden_dim) + return final_hidden_states, router_logits + + +# --------------------------------------------------------------------------- +# Attention with partial rotary embeddings +# --------------------------------------------------------------------------- + + +class Glm4MoeAttention(MultiHeadAttention): + """GLM-4 MoE attention with partial rotary position embeddings. + + GLM uses ``partial_rotary_factor=0.5``: only the first half of each head + dimension receives RoPE. The ``o_proj`` never has bias. + """ + + def __init__(self, config, layer_idx: int): + super().__init__(config, layer_idx) + # GLM's o_proj has no bias regardless of attention_bias setting + if self.o_proj.bias is not None: + self.o_proj = nn.Linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=False) + # QK norm is controlled by config + if not getattr(config, "use_qk_norm", False): + self.q_norm = nn.Identity() + self.k_norm = nn.Identity() + + def _project_qkv( + self, + hidden_states: torch.Tensor, + position_embeddings: Tuple[torch.Tensor, torch.Tensor], + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + input_shape = hidden_states.shape[:-1] + hidden_shape = (*input_shape, -1, self.head_dim) + + if hasattr(self, "qkv_proj"): + qkv = self.qkv_proj(hidden_states) + q, k, v = qkv.split([self.q_dim, self.kv_dim, self.kv_dim], dim=-1) + else: + 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)) + v = v.view(hidden_shape) + + # Partial rotary: only apply RoPE to the first `rotary_dim` dimensions + cos, sin = position_embeddings + rotary_dim = cos.shape[-1] + q_rot, q_pass = q[..., :rotary_dim], q[..., rotary_dim:] + k_rot, k_pass = k[..., :rotary_dim], k[..., rotary_dim:] + q_rot, k_rot = apply_rotary_pos_emb(q_rot, k_rot, cos, sin) + q = torch.cat([q_rot, q_pass], dim=-1) + k = torch.cat([k_rot, k_pass], dim=-1) + + if getattr(self.config, "_attention_cast_bf16", False): + q = q.to(torch.bfloat16) + k = k.to(torch.bfloat16) + + return q, k, v + + +# --------------------------------------------------------------------------- +# Decoder layer +# --------------------------------------------------------------------------- + +GLM4_MOE_CLASSES = { + "eager": lambda config: Glm4MoeSparseMoeBlock(config, moe_implementation="eager"), + "triton": lambda config: Glm4MoeSparseMoeBlock(config, moe_implementation="triton"), + "native": lambda config: Glm4MoeSparseMoeBlock(config, moe_implementation="native"), + "quack": lambda config: Glm4MoeSparseMoeBlock(config, moe_implementation="quack"), +} + + +class Glm4MoeDecoderLayer(nn.Module): + def __init__(self, config: Glm4MoeConfig, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + + self.self_attn = Glm4MoeAttention(config, layer_idx) + + self.input_layernorm = Glm4MoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = Glm4MoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + if layer_idx >= config.first_k_dense_replace: + moe_implementation = getattr(config, "_moe_implementation", "triton") + self.mlp = GLM4_MOE_CLASSES[moe_implementation](config) + else: + self.mlp = Glm4MoeMLP(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, ...]: + _selective = ( + self.training + and getattr(self, "gradient_checkpointing", False) + and getattr(self, "_recompute_modules", None) is not None + ) + residual = hidden_states + + hidden_states = self.input_layernorm(hidden_states) + + if _selective and "self_attn" in self._recompute_modules: + 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 + + 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_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 + + +# --------------------------------------------------------------------------- +# PreTrainedModel / Model / ForCausalLM +# --------------------------------------------------------------------------- + + +class Glm4MoePreTrainedModel(XorlPreTrainedModel): + config_class = Glm4MoeConfig + base_model_prefix = "model" + _no_split_modules = ["Glm4MoeDecoderLayer"] + + 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, Glm4MoeRMSNorm): + module.weight.data.fill_(1.0) + elif isinstance(module, Glm4MoeGate): + nn.init.kaiming_uniform_(module.weight) + module.e_score_correction_bias.data.zero_() + 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_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) + + 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 + ) + + 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) + return Glm4MoeCheckpointHandler( + num_experts=self.config.n_routed_experts, + num_attention_heads=self.config.num_attention_heads, + num_key_value_heads=self.config.num_key_value_heads, + head_dim=head_dim, + ep_rank=ep_rank, + ep_size=ep_size, + checkpoint_has_per_expert=has_per_expert, + skip_qkv_merge=unfused, + skip_gate_up_merge=unfused, + is_prequantized=is_prequantized, + exclude_modules=exclude_modules, + device=kwargs.get("device"), + model=self if is_prequantized else None, + num_hidden_layers=self.config.num_hidden_layers, + ) + + +class Glm4MoeModel(Glm4MoePreTrainedModel): + def __init__(self, config: Glm4MoeConfig): + 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( + [Glm4MoeDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self.norm = Glm4MoeRMSNorm(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: 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, + output_router_logits: Optional[bool] = None, + **kwargs: Unpack[AttentionKwargs], + ) -> MoeModelOutput: + output_attentions = ( + output_attentions if output_attentions is not None else getattr(self.config, "output_attentions", False) + ) + output_router_logits = ( + output_router_logits + if output_router_logits is not None + else getattr(self.config, "output_router_logits", False) + ) + output_hidden_states = ( + output_hidden_states + if output_hidden_states is not None + else getattr(self.config, "output_hidden_states", False) + ) + + 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, + sliding_window=getattr(self.config, "sliding_window", None), + is_training=self.training, + output_attentions=output_attentions, + ) + + 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 + all_router_logits = () if output_router_logits else None + + for decoder_layer in self.layers: + if decoder_layer is None: + continue + _use_outer_checkpoint = ( + self.gradient_checkpointing and self.training and getattr(self, "_recompute_modules", None) is None + ) + + if _use_outer_checkpoint: + layer_outputs = self._gradient_checkpointing_func( + decoder_layer.__call__, + hidden_states, + causal_mask, + position_ids, + output_attentions, + output_router_logits, + position_embeddings, + **kwargs, + ) + else: + 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, + ) + + hidden_states = layer_outputs[0] + + if output_attentions: + all_self_attns += (layer_outputs[1],) + + 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, + ) + + +class KwargsForCausalLM(AttentionKwargs): ... + + +class Glm4MoeForCausalLM(Glm4MoePreTrainedModel): + _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 = Glm4MoeModel(config) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.router_aux_loss_coef = getattr(config, "router_aux_loss_coef", 0.001) + self.num_experts = config.n_routed_experts + self.num_experts_per_tok = config.num_experts_per_tok + + self.post_init() + + def unfuse_for_tp(self): + 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: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + **kwargs, + ) -> MoeCausalLMOutput: + output_router_logits = getattr(self.config, "output_router_logits", False) + + outputs = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + output_router_logits=output_router_logits, + **kwargs, + ) + + return MoeCausalLMOutput( + last_hidden_state=outputs.last_hidden_state, + router_logits=outputs.router_logits, + ) + + +ModelClass = Glm4MoeForCausalLM + + +__all__ = [ + "Glm4MoeForCausalLM", + "Glm4MoeModel", + "Glm4MoePreTrainedModel", + "Glm4MoeSparseMoeBlock", + "Glm4MoeGate", + "Glm4MoeMLP", + "Glm4MoeAttention", +] diff --git a/src/xorl/models/transformers/glm4_moe/parallelize.py b/src/xorl/models/transformers/glm4_moe/parallelize.py new file mode 100644 index 00000000..e998f663 --- /dev/null +++ b/src/xorl/models/transformers/glm4_moe/parallelize.py @@ -0,0 +1,63 @@ +"""Parallelization plan and utilities for GLM-4 MoE models.""" + +from torch.distributed._tensor import Shard + +from ....distributed.parallel_plan import ParallelPlan +from ...layers.moe import MoEBlock + + +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", + # Dense MLP layers (first K layers before MoE) + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise", + # Shared expert MLP (inside MoE blocks) + "layers.*.mlp.shared_experts.gate_proj": "colwise", + "layers.*.mlp.shared_experts.up_proj": "colwise", + "layers.*.mlp.shared_experts.down_proj": "rowwise", +} + +MODEL_TP_PLAN = { + "lm_head": "colwise_rep", +} + + +def unfuse_for_tp(model): + """Unfuse fused projections for tensor parallelism compatibility. + + For ALL layers: splits ``qkv_proj`` -> ``q_proj / k_proj / v_proj`` in attention. + For DENSE layers only: splits ``gate_up_proj`` -> ``gate_proj / up_proj`` in MLP. + For MoE layers: also unfuses the shared expert MLP. + MoE expert weights are not TP-sharded (they use EP instead). + """ + for layer in model.model.layers: + layer.self_attn.unfuse_for_tp() + if isinstance(layer.mlp, MoEBlock): + if hasattr(layer.mlp, "shared_experts") and layer.mlp.shared_experts is not None: + layer.mlp.shared_experts.unfuse_for_tp() + else: + layer.mlp.unfuse_for_tp() + model._unfused_for_tp = True + model.config.base_model_tp_plan = TP_PLAN + + +def get_ep_plan(): + """Get EP (expert parallelism) plan for GLM-4 MoE model.""" + ep_plan = { + # Expert base weights — stored as fused gate_up_proj [E, H, 2*I] + "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), + "model.layers.*.mlp.experts.gate_proj_lora_B": Shard(0), + "model.layers.*.mlp.experts.up_proj_lora_A": Shard(0), + "model.layers.*.mlp.experts.up_proj_lora_B": Shard(0), + "model.layers.*.mlp.experts.down_proj_lora_A": Shard(0), + "model.layers.*.mlp.experts.down_proj_lora_B": Shard(0), + } + return ParallelPlan(ep_plan=ep_plan) diff --git a/src/xorl/models/transformers/gpt_oss/__init__.py b/src/xorl/models/transformers/gpt_oss/__init__.py new file mode 100644 index 00000000..4dabd8a6 --- /dev/null +++ b/src/xorl/models/transformers/gpt_oss/__init__.py @@ -0,0 +1,8 @@ +"""GPT-OSS model.""" + +from .configuration_gpt_oss import GptOssConfig + + +__all__ = [ + "GptOssConfig", +] diff --git a/src/xorl/models/transformers/gpt_oss/checkpoint_handler.py b/src/xorl/models/transformers/gpt_oss/checkpoint_handler.py new file mode 100644 index 00000000..46eed7f8 --- /dev/null +++ b/src/xorl/models/transformers/gpt_oss/checkpoint_handler.py @@ -0,0 +1,410 @@ +"""Checkpoint handler for GPT-OSS models. + +GPT-OSS original checkpoints use non-standard key names. Expert weights are +stored in MXFP4 quantized format (``.blocks`` + ``.scales`` pairs) and need +dequantization during loading. Non-expert weights (attention, norms, +embeddings) are BF16 and just need key renaming. + +Load transforms: + 1. Rename keys → xorl internal parameter names + 2. Buffer MXFP4 ``.blocks``/``.scales`` pairs, dequantize to BF16 + 3. Transpose expert weights from [E, out, in] → [E, in, out] + 4. Deinterleave HF gate/up tensors when needed + 5. Slice expert dimension for EP when ep_size > 1 + +Save transforms: + 1. Preserve xorl's normalized parameter names so saved checkpoints remain + directly reloadable by xorl. +""" + +import re +import warnings +from typing import Dict, List, Optional, Tuple + +import torch +from torch import Tensor + +from ...checkpoint_handlers.base import CheckpointHandler + + +# ============================================================================ +# MXFP4 dequantization (CPU-compatible, no Triton dependency) +# ============================================================================ + +# FP4 E2M1 lookup table: 4-bit code → float32 value +# Codes 0-7 are positive, 8-15 are negative (sign bit = bit 3) +_FP4_E2M1_LUT = torch.tensor( + [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, 0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0], + dtype=torch.float32, +) + + +def _mxfp4_dequantize_cpu( + blocks: Tensor, + scales: Tensor, +) -> Tensor: + """Dequantize MXFP4 packed expert weights to bfloat16. + + Args: + blocks: uint8 tensor of shape [..., num_groups, 16] (packed FP4 pairs) + scales: uint8 tensor of shape [..., num_groups] (E8M0 exponents) + + Returns: + bfloat16 tensor of shape [..., num_groups * 32] (32 values per group) + """ + # Unpack two FP4 values per byte: low nibble = even index, high nibble = odd + lo = (blocks & 0x0F).to(torch.int64) + hi = (blocks >> 4).to(torch.int64) + + # Lookup float values: lo/hi each [..., num_groups, 16] + lut = _FP4_E2M1_LUT.to(blocks.device) + val_lo = lut[lo] # [..., G, 16] + val_hi = lut[hi] # [..., G, 16] + + # Interleave: even positions = lo, odd positions = hi → [..., G, 32] + interleaved = torch.stack([val_lo, val_hi], dim=-1) # [..., G, 16, 2] + interleaved = interleaved.reshape(*blocks.shape[:-1], 32) # [..., G, 32] + + # E8M0 scale: 2^(scale - 127) + scale_f32 = torch.pow(2.0, scales.to(torch.float32) - 127.0) # [..., G] + + # Broadcast multiply: [..., G, 1] * [..., G, 32] + result = interleaved * scale_f32.unsqueeze(-1) + + # Flatten groups: [..., G, 32] → [..., G * 32] + result = result.reshape(*scales.shape[:-1], -1) + + return result.to(torch.bfloat16) + + +# ============================================================================ +# Key mapping: original checkpoint → xorl internal +# ============================================================================ + +_LOAD_KEY_MAP = [ + # Embeddings + (re.compile(r"^embedding\.weight$"), "model.embed_tokens.weight", False), + (re.compile(r"^unembedding\.weight$"), "lm_head.weight", False), + # Final norm + (re.compile(r"^norm\.scale$"), "model.norm.weight", False), + # Attention norm + (re.compile(r"^block\.(\d+)\.attn\.norm\.scale$"), r"model.layers.\1.input_layernorm.weight", False), + # Attention output + (re.compile(r"^block\.(\d+)\.attn\.out\.(weight|bias)$"), r"model.layers.\1.self_attn.o_proj.\2", False), + # Attention sinks + (re.compile(r"^block\.(\d+)\.attn\.sinks$"), r"model.layers.\1.self_attn.sinks", False), + # MLP norm + (re.compile(r"^block\.(\d+)\.mlp\.norm\.scale$"), r"model.layers.\1.post_attention_layernorm.weight", False), + # Router gate + (re.compile(r"^block\.(\d+)\.mlp\.gate\.(weight|bias)$"), r"model.layers.\1.mlp.gate.\2", False), + # Expert biases (BF16, no transpose needed) + (re.compile(r"^block\.(\d+)\.mlp\.mlp1_bias$"), r"model.layers.\1.mlp.experts.gate_up_bias", False), + (re.compile(r"^block\.(\d+)\.mlp\.mlp2_bias$"), r"model.layers.\1.mlp.experts.down_bias", False), +] + +# Pattern for MXFP4 expert weight keys (original checkpoint format) +_MXFP4_PATTERN = re.compile(r"^block\.(\d+)\.mlp\.(mlp[12])_weight\.(blocks|scales)$") + +# Internal names for expert weights after dequant +_EXPERT_INTERNAL_NAMES = { + "mlp1": "mlp.experts.gate_up_proj", + "mlp2": "mlp.experts.down_proj", +} + + +# ============================================================================ +# HF checkpoint format support +# ============================================================================ + +# HF format key renames → xorl internal +_HF_LOAD_KEY_MAP = [ + # Router → gate + (re.compile(r"^model\.layers\.(\d+)\.mlp\.router\.(weight|bias)$"), r"model.layers.\1.mlp.gate.\2", False), + # Expert biases (gate_up_proj_bias → gate_up_bias, drop _proj) + ( + re.compile(r"^model\.layers\.(\d+)\.mlp\.experts\.gate_up_proj_bias$"), + r"model.layers.\1.mlp.experts.gate_up_bias", + False, + ), + ( + re.compile(r"^model\.layers\.(\d+)\.mlp\.experts\.down_proj_bias$"), + r"model.layers.\1.mlp.experts.down_bias", + False, + ), +] + +# HF format MXFP4: gate_up_proj_blocks/_scales, down_proj_blocks/_scales +_HF_MXFP4_PATTERN = re.compile(r"^model\.layers\.(\d+)\.mlp\.experts\.(gate_up|down)_proj_(blocks|scales)$") + +_HF_EXPERT_INTERNAL_NAMES = { + "gate_up": "mlp.experts.gate_up_proj", + "down": "mlp.experts.down_proj", +} + +# HF format separate q/k/v projections that need fusing +_HF_QKV_PATTERN = re.compile(r"^model\.layers\.(\d+)\.self_attn\.(q|k|v)_proj\.(weight|bias)$") + +# Original and xorl-local fused qkv projections +_ORIGINAL_FUSED_QKV_PATTERN = re.compile(r"^block\.(\d+)\.attn\.qkv\.(weight|bias)$") +_INTERNAL_FUSED_QKV_PATTERN = re.compile(r"^model\.layers\.(\d+)\.self_attn\.qkv_proj\.(weight|bias)$") + +# Keys holding stacked expert tensors (need EP slicing) +_EXPERT_KEY_PATTERN = re.compile(r"^model\.layers\.\d+\.mlp\.experts\.(gate_up_proj|gate_up_bias|down_proj|down_bias)$") + + +def _detect_checkpoint_format(checkpoint_keys: Optional[set[str]]) -> str: + """Best-effort source format detection for ambiguous GPT-OSS keys.""" + if not checkpoint_keys: + return "xorl" + + if any( + key.startswith("block.") or key in {"embedding.weight", "unembedding.weight", "norm.scale"} + for key in checkpoint_keys + ): + return "original" + + if any( + ".mlp.gate." in key + or ".self_attn.qkv_proj." in key + or key.endswith(".mlp.experts.gate_up_bias") + or key.endswith(".mlp.experts.down_bias") + for key in checkpoint_keys + ): + return "xorl" + + if any( + ".mlp.router." in key + or key.endswith(".mlp.experts.gate_up_proj_bias") + or key.endswith(".mlp.experts.down_proj_bias") + for key in checkpoint_keys + ): + return "hf" + + return "xorl" + + +def _remap_key(key: str, key_map, tensor: Tensor) -> Optional[Tuple[str, Tensor]]: + """Apply the first matching pattern from *key_map*.""" + for pattern, replacement, transpose in key_map: + new_key, n = pattern.subn(replacement, key) + if n > 0: + if transpose and tensor.ndim == 3: + tensor = tensor.transpose(1, 2).contiguous() + return new_key, tensor + return None + + +def _deinterleave_gate_up(tensor: Tensor) -> Tensor: + """Convert interleaved ``[g0,u0,g1,u1,...]`` to ``[g0,g1,...|u0,u1,...]`` in the last dim. + + The original GPT-OSS checkpoint stores gate and up values interleaved. + The standard xorl MoE layout concatenates gate then up. + """ + gate = tensor[..., 0::2] + up = tensor[..., 1::2] + return torch.cat([gate, up], dim=-1).contiguous() + + +def _split_fused_qkv( + tensor: Tensor, + *, + layer_idx: str, + wb: str, + q_dim: int, + kv_dim: int, +) -> List[Tuple[str, Tensor]]: + """Split fused qkv tensor into q/k/v tensors under xorl/HF-style names.""" + prefix = f"model.layers.{layer_idx}.self_attn" + q = tensor[:q_dim].contiguous() + k = tensor[q_dim : q_dim + kv_dim].contiguous() + v = tensor[q_dim + kv_dim : q_dim + (2 * kv_dim)].contiguous() + return [ + (f"{prefix}.q_proj.{wb}", q), + (f"{prefix}.k_proj.{wb}", k), + (f"{prefix}.v_proj.{wb}", v), + ] + + +class GptOssCheckpointHandler(CheckpointHandler): + """Checkpoint handler for GPT-OSS models. + + Handles original GPT-OSS checkpoints, HF-style GPT-OSS checkpoints, and + xorl-local normalized checkpoints. + + Args: + num_experts: Total number of experts. + num_attention_heads: Number of query heads (used to split fused qkv). + num_key_value_heads: Number of key/value heads (used to split fused qkv). + head_dim: Attention head dimension. + ep_rank: This rank's expert-parallel index (default 0). + ep_size: Total number of expert-parallel ranks (default 1). + checkpoint_keys: Optional full checkpoint key set for source-format detection. + skip_qkv_merge: When True, emit separate q/k/v tensors instead of fused + ``qkv_proj`` tensors (used after ``merge_qkv=False`` / TP unfuse). + """ + + def __init__( + self, + num_experts: int, + num_attention_heads: int, + num_key_value_heads: int, + head_dim: int, + ep_rank: int = 0, + ep_size: int = 1, + checkpoint_keys: Optional[set[str]] = None, + skip_qkv_merge: bool = False, + ): + self._num_experts = num_experts + 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._q_dim = num_attention_heads * head_dim + self._kv_dim = num_key_value_heads * head_dim + self._skip_qkv_merge = skip_qkv_merge + self._checkpoint_format = _detect_checkpoint_format(checkpoint_keys) + + # Buffer for MXFP4 pairs: {(layer_idx, mlp_name): {"blocks": ..., "scales": ...}} + self._mxfp4_buffer: Dict[Tuple[str, str], Dict[str, Tensor]] = {} + # Buffer for HF separate q/k/v → fused qkv: {(layer_idx, "weight"|"bias"): {"q": ..., "k": ..., "v": ...}} + self._qkv_buffer: Dict[Tuple[str, str], Dict[str, Tensor]] = {} + + def _slice_expert_tensor_for_ep(self, key: str, tensor: Tensor) -> Tensor: + """Slice stacked expert tensor along dim 0 for expert parallelism.""" + if self._ep_size <= 1: + return tensor + if _EXPERT_KEY_PATTERN.match(key): + return tensor[self._expert_start : self._expert_end].contiguous() + return tensor + + def _handle_mxfp4_dequant( + self, + layer_idx: str, + mlp_name: str, + part: str, + tensor: Tensor, + internal_names: dict, + ) -> List[Tuple[str, Tensor]]: + """Buffer MXFP4 blocks/scales and dequantize when both arrive.""" + buf_key = (layer_idx, mlp_name) + + if buf_key not in self._mxfp4_buffer: + self._mxfp4_buffer[buf_key] = {} + self._mxfp4_buffer[buf_key][part] = tensor + + if len(self._mxfp4_buffer[buf_key]) < 2: + return [] + + parts = self._mxfp4_buffer.pop(buf_key) + weight_bf16 = _mxfp4_dequantize_cpu(parts["blocks"], parts["scales"]) + # [E, out_dim, in_dim] → [E, in_dim, out_dim] + weight_bf16 = weight_bf16.transpose(1, 2).contiguous() + + # Deinterleave gate_up from [g0,u0,g1,u1,...] to [g0,g1,...|u0,u1,...]. + if mlp_name in ("mlp1", "gate_up"): + weight_bf16 = _deinterleave_gate_up(weight_bf16) + + internal_name = f"model.layers.{layer_idx}.{internal_names[mlp_name]}" + weight_bf16 = self._slice_expert_tensor_for_ep(internal_name, weight_bf16) + return [(internal_name, weight_bf16)] + + def on_load_weight(self, key: str, tensor: Tensor) -> List[Tuple[str, Tensor]]: + # --- Original format: MXFP4 expert weights --- + m = _MXFP4_PATTERN.match(key) + if m is not None: + return self._handle_mxfp4_dequant( + m.group(1), + m.group(2), + m.group(3), + tensor, + _EXPERT_INTERNAL_NAMES, + ) + + # --- HF format: MXFP4 expert weights --- + m = _HF_MXFP4_PATTERN.match(key) + if m is not None: + return self._handle_mxfp4_dequant( + m.group(1), + m.group(2), + m.group(3), + tensor, + _HF_EXPERT_INTERNAL_NAMES, + ) + + # --- Original format: fused qkv --- + m = _ORIGINAL_FUSED_QKV_PATTERN.match(key) + if m is not None: + layer_idx, wb = m.group(1), m.group(2) + if self._skip_qkv_merge: + return _split_fused_qkv(tensor, layer_idx=layer_idx, wb=wb, q_dim=self._q_dim, kv_dim=self._kv_dim) + internal_name = f"model.layers.{layer_idx}.self_attn.qkv_proj.{wb}" + return [(internal_name, tensor)] + + # --- xorl-local format: fused qkv --- + m = _INTERNAL_FUSED_QKV_PATTERN.match(key) + if m is not None: + layer_idx, wb = m.group(1), m.group(2) + if self._skip_qkv_merge: + return _split_fused_qkv(tensor, layer_idx=layer_idx, wb=wb, q_dim=self._q_dim, kv_dim=self._kv_dim) + return [(key, tensor)] + + # --- HF/xorl-unfused format: separate q/k/v --- + m = _HF_QKV_PATTERN.match(key) + if m is not None: + if self._skip_qkv_merge: + return [(key, tensor)] + + layer_idx, qkv, wb = m.group(1), m.group(2), m.group(3) + buf_key = (layer_idx, wb) + + if buf_key not in self._qkv_buffer: + self._qkv_buffer[buf_key] = {} + self._qkv_buffer[buf_key][qkv] = tensor + + if len(self._qkv_buffer[buf_key]) < 3: + return [] + + parts = self._qkv_buffer.pop(buf_key) + fused = torch.cat([parts["q"], parts["k"], parts["v"]], dim=0) + internal_name = f"model.layers.{layer_idx}.self_attn.qkv_proj.{wb}" + return [(internal_name, fused)] + + # --- xorl/HF stacked expert tensors --- + if _EXPERT_KEY_PATTERN.match(key): + if key.endswith("gate_up_proj") and self._checkpoint_format == "hf": + tensor = _deinterleave_gate_up(tensor) + tensor = self._slice_expert_tensor_for_ep(key, tensor) + return [(key, tensor)] + + # --- HF format: simple renames (router→gate, expert biases) --- + result = _remap_key(key, _HF_LOAD_KEY_MAP, tensor) + if result is not None: + new_key, tensor = result + if "gate_up_bias" in new_key: + tensor = _deinterleave_gate_up(tensor) + tensor = self._slice_expert_tensor_for_ep(new_key, tensor) + return [(new_key, tensor)] + + # --- Original format: standard key remapping --- + result = _remap_key(key, _LOAD_KEY_MAP, tensor) + if result is not None: + new_key, tensor = result + if "gate_up_bias" in new_key: + tensor = _deinterleave_gate_up(tensor) + tensor = self._slice_expert_tensor_for_ep(new_key, tensor) + return [(new_key, tensor)] + + # Pass-through (HF keys that already match xorl internal names) + return [(key, tensor)] + + def on_load_complete(self) -> List[Tuple[str, Tensor]]: + if self._mxfp4_buffer: + incomplete = [f"layer {l} {m}" for (l, m) in self._mxfp4_buffer.keys()] + warnings.warn(f"Incomplete MXFP4 pairs at load completion: {incomplete}") + self._mxfp4_buffer.clear() + if self._qkv_buffer: + incomplete = [f"layer {l} qkv {wb}" for (l, wb) in self._qkv_buffer.keys()] + warnings.warn(f"Incomplete QKV groups at load completion: {incomplete}") + self._qkv_buffer.clear() + return [] diff --git a/src/xorl/models/transformers/gpt_oss/configuration_gpt_oss.py b/src/xorl/models/transformers/gpt_oss/configuration_gpt_oss.py new file mode 100644 index 00000000..288831b4 --- /dev/null +++ b/src/xorl/models/transformers/gpt_oss/configuration_gpt_oss.py @@ -0,0 +1,204 @@ +"""GPT-OSS model configuration.""" + +from transformers.configuration_utils import PretrainedConfig + +from .parallelize import TP_PLAN + + +class GptOssConfig(PretrainedConfig): + r""" + Configuration class for GPT-OSS models (e.g. ``openai/gpt-oss-20b``). + + GPT-OSS is a Mixture-of-Experts transformer with: + - SwiGLU activation with clamping (``swiglu_limit``) + - Attention sinks (learned per-head bias added to softmax) + - Alternating sliding-window / full attention layers + - Custom YaRN-style RoPE with NTK-by-parts scaling + - Expert biases on both gate_up (mlp1) and down (mlp2) projections + - Router gate with bias + + Args: + vocab_size (``int``, defaults to 201088): + Vocabulary size. + hidden_size (``int``, defaults to 2880): + Hidden dimension. + moe_intermediate_size (``int``, defaults to 2880): + Expert FFN intermediate dimension. + num_hidden_layers (``int``, defaults to 24): + Number of transformer layers. + num_attention_heads (``int``, defaults to 64): + Number of query attention heads. + num_key_value_heads (``int``, defaults to 8): + Number of key/value heads for GQA. + head_dim (``int``, defaults to 64): + Dimension per attention head. + rms_norm_eps (``float``, defaults to 1e-5): + Epsilon for RMS normalization. + num_experts (``int``, defaults to 32): + Total number of MoE experts. + num_experts_per_tok (``int``, defaults to 4): + Number of experts activated per token. + norm_topk_prob (``bool``, defaults to True): + Whether to renormalize top-k routing weights. + swiglu_limit (``float``, defaults to 7.0): + Clamp limit for SwiGLU activation inputs. + hidden_act (``str``, defaults to ``"silu"``): + Base activation function name. + attention_bias (``bool``, defaults to True): + Whether QKV and output projections have bias terms. + attention_dropout (``float``, defaults to 0.0): + Dropout ratio for attention weights. + sliding_window (``int``, defaults to 128): + Sliding window size (applied to even-indexed layers). + rope_theta (``float``, defaults to 150000.0): + Base period for RoPE embeddings. + initial_context_length (``int``, defaults to 4096): + Original pre-training context length (used by YaRN RoPE). + rope_scaling_factor (``float``, defaults to 32.0): + Context extension scaling factor for YaRN RoPE. + rope_ntk_alpha (``float``, defaults to 1.0): + NTK-by-parts alpha (high-frequency boundary). + rope_ntk_beta (``float``, defaults to 32.0): + NTK-by-parts beta (low-frequency boundary). + max_position_embeddings (``int``, defaults to 131072): + Maximum sequence length the model can handle. + """ + + model_type = "gpt_oss" + + base_model_tp_plan = TP_PLAN + + def __init__( + self, + vocab_size=201088, + hidden_size=2880, + moe_intermediate_size=2880, + num_hidden_layers=24, + num_attention_heads=64, + num_key_value_heads=8, + head_dim=64, + rms_norm_eps=1e-5, + use_cache=False, + tie_word_embeddings=False, + # MoE + num_experts=32, + num_experts_per_tok=4, + norm_topk_prob=True, + # Activation + swiglu_limit=7.0, + hidden_act="silu", + # Attention + attention_bias=True, + attention_dropout=0.0, + sliding_window=128, + # RoPE + rope_theta=150000.0, + initial_context_length=4096, + rope_scaling_factor=32.0, + rope_ntk_alpha=1.0, + rope_ntk_beta=32.0, + max_position_embeddings=131072, + # Internal + _moe_implementation="native", + **kwargs, + ): + self.vocab_size = vocab_size + self.hidden_size = hidden_size + self.moe_intermediate_size = moe_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.rms_norm_eps = rms_norm_eps + self.use_cache = use_cache + self.max_position_embeddings = max_position_embeddings + + # MoE + self.num_experts = num_experts + self.num_experts_per_tok = num_experts_per_tok + self.norm_topk_prob = norm_topk_prob + self._moe_implementation = _moe_implementation + + # Activation + self.swiglu_limit = swiglu_limit + self.hidden_act = hidden_act + + # Attention + self.attention_bias = attention_bias + self.attention_dropout = attention_dropout + self.sliding_window = sliding_window + + # RoPE — custom YaRN with NTK-by-parts + self.rope_theta = rope_theta + self.initial_context_length = initial_context_length + self.rope_scaling_factor = rope_scaling_factor + self.rope_ntk_alpha = rope_ntk_alpha + self.rope_ntk_beta = rope_ntk_beta + + super().__init__( + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) + + @classmethod + @staticmethod + def _resolve_rope_params(hf_config): + """Extract RoPE params from either flat fields or HF nested rope_scaling dict.""" + rope_scaling = getattr(hf_config, "rope_scaling", None) + if isinstance(rope_scaling, dict): + return dict( + initial_context_length=rope_scaling.get("original_max_position_embeddings", 4096), + rope_scaling_factor=rope_scaling.get("factor", 32.0), + rope_ntk_alpha=rope_scaling.get("beta_slow", 1.0), + rope_ntk_beta=rope_scaling.get("beta_fast", 32.0), + ) + return dict( + initial_context_length=getattr(hf_config, "initial_context_length", 4096), + rope_scaling_factor=getattr(hf_config, "rope_scaling_factor", 32.0), + rope_ntk_alpha=getattr(hf_config, "rope_ntk_alpha", 1.0), + rope_ntk_beta=getattr(hf_config, "rope_ntk_beta", 32.0), + ) + + @classmethod + def from_hf_config(cls, hf_config): + """Build a GptOssConfig from an HF config dict/namespace. + + The GPT-OSS HF config uses non-standard field names (e.g. + ``experts_per_token`` instead of ``num_experts_per_tok``, + ``intermediate_size`` for the expert FFN dimension). + """ + return cls( + vocab_size=getattr(hf_config, "vocab_size", 201088), + hidden_size=getattr(hf_config, "hidden_size", 2880), + moe_intermediate_size=getattr(hf_config, "moe_intermediate_size", None) + or getattr(hf_config, "intermediate_size", 2880), + num_hidden_layers=getattr(hf_config, "num_hidden_layers", 24), + num_attention_heads=getattr(hf_config, "num_attention_heads", 64), + num_key_value_heads=getattr(hf_config, "num_key_value_heads", 8), + head_dim=getattr(hf_config, "head_dim", 64), + rms_norm_eps=getattr(hf_config, "rms_norm_eps", 1e-5), + use_cache=getattr(hf_config, "use_cache", False), + # HF uses "num_local_experts", original uses "num_experts" + num_experts=getattr(hf_config, "num_experts", None) or getattr(hf_config, "num_local_experts", 32), + # HF uses "num_experts_per_tok", original uses "experts_per_token" + num_experts_per_tok=getattr(hf_config, "experts_per_token", None) + or getattr(hf_config, "num_experts_per_tok", 4), + swiglu_limit=getattr(hf_config, "swiglu_limit", 7.0), + hidden_act=getattr(hf_config, "hidden_act", "silu"), + attention_dropout=getattr(hf_config, "attention_dropout", 0.0), + sliding_window=getattr(hf_config, "sliding_window", 128), + rope_theta=getattr(hf_config, "rope_theta", 150000.0), + max_position_embeddings=getattr(hf_config, "max_position_embeddings", 131072), + # Support both flat fields (original) and nested rope_scaling dict (HF). + # The HF config stores rope params in a nested dict; the original uses flat fields. + **cls._resolve_rope_params(hf_config), + # GPT-OSS always has attention bias (nn.Linear defaults) + attention_bias=getattr(hf_config, "attention_bias", True), + tie_word_embeddings=getattr(hf_config, "tie_word_embeddings", False), + _moe_implementation=getattr(hf_config, "_moe_implementation", "native"), + architectures=getattr(hf_config, "architectures", ["GptOssForCausalLM"]), + output_router_logits=getattr(hf_config, "output_router_logits", False), + ) + + +__all__ = ["GptOssConfig"] diff --git a/src/xorl/models/transformers/gpt_oss/flash_sink_attention.py b/src/xorl/models/transformers/gpt_oss/flash_sink_attention.py new file mode 100644 index 00000000..81c408a9 --- /dev/null +++ b/src/xorl/models/transformers/gpt_oss/flash_sink_attention.py @@ -0,0 +1,315 @@ +"""Sink-aware Flash Attention 3 for GPT-OSS. + +Wraps FA3's forward with an autograd function that fuses the learned +per-head sink into the softmax via post-hoc multiplication by +``sigmoid(lse - sink)``. Mathematically equivalent to eager's concat-sink +-then-softmax, just using FA3's fused kernel for the underlying attention. + +Needed because the installed ``flash_attn_interface`` build does not yet +expose a native ``sinks=`` kwarg; when that kernel support lands, this file +can be replaced by a direct passthrough. + +A single autograd function handles both batched and packed-varlen inputs: +the only layout difference the sink math cares about is the rank of the +``lse`` tensor (``[B, Hq, S]`` batched vs ``[Hq, T]`` varlen), which broadcasts +uniformly once we keep *head* on axis ``-2`` and *sequence* on axis ``-1``. + +Backward decomposition +---------------------- +Let ``m = sigmoid(lse - sink)``. Then ``o = o_flash * m``, and gradients +decompose into four paths: + + (1) main: dq, dk, dv via FA backward with ``dout' = dO * m`` + (2) dsink: -sum(g_r * m * (1-m)) where g_r = (dO * o_flash).sum(-1) + (3) dq extra: scale * g_ell * attention(Q, K, K) (extra FA fwd) + (4) dk extra: scale * P^T (g_ell * Q) (FA bwd, dv slot) + +where ``g_ell = g_r * m * (1-m)``. The derivation uses +``dlse/dq = scale * PK`` and ``dlse/dk = scale * P^T (scalar * Q)``. +""" + +import math +from typing import Optional, Tuple + +import torch +from flash_attn_interface import ( + _flash_attn_backward, + flash_attn_func, + flash_attn_varlen_func, +) + + +def _fa3_forward( + q, + k, + v, + *, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + softmax_scale, + causal, + window_size, + softcap, + deterministic, + return_attn_probs, +): + common = dict( + softmax_scale=softmax_scale, + causal=causal, + window_size=window_size, + softcap=softcap, + deterministic=deterministic, + return_attn_probs=return_attn_probs, + ) + if cu_seqlens_q is not None: + return flash_attn_varlen_func( + q, + k, + v, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + **common, + ) + return flash_attn_func(q, k, v, **common) + + +def _fa3_backward( + dout, + q, + k, + v, + out, + lse, + dq, + dk, + dv, + *, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + softmax_scale, + causal, + window_size, + softcap, + deterministic, +): + kwargs = dict( + dq=dq, + dk=dk, + dv=dv, + softmax_scale=softmax_scale, + is_causal=causal, + window_size_left=window_size[0], + window_size_right=window_size[1], + softcap=softcap, + deterministic=deterministic, + ) + if cu_seqlens_q is not None: + 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(dout, q, k, v, out, lse, **kwargs) + + +class FlashAttnWithSinkFA3(torch.autograd.Function): + """FA3 attention with a learned per-head sink in the softmax denominator. + + Handles batched (4D q/k/v, 3D lse) and packed-varlen (3D q/k/v, 2D lse) + uniformly by requiring that in both layouts ``lse`` has head on axis -2 + and sequence on axis -1, and ``out`` has sequence, head, dim as the last + three axes. + """ + + @staticmethod + def forward( + ctx, + q, + k, + v, + sink, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + causal, + window_size, + softmax_scale, + softcap, + deterministic, + ): + out, lse = _fa3_forward( + q, + k, + v, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + softmax_scale=softmax_scale, + causal=causal, + window_size=window_size, + softcap=softcap, + deterministic=deterministic, + return_attn_probs=True, + ) + # Broadcast sink over lse: lse is [..., Hq, S] (head on -2, seq on -1). + sink_f = sink.float().view(*([1] * (lse.ndim - 2)), -1, 1) + m = torch.sigmoid(lse - sink_f) + # Align m to out's layout ([..., S, Hq, D]): move seq next to head-dim slot. + m_for_out = m.transpose(-2, -1).unsqueeze(-1).to(out.dtype) + out_final = out * m_for_out + + saved = [q, k, v, sink, out, lse] + if cu_seqlens_q is not None: + saved.extend([cu_seqlens_q, cu_seqlens_k]) + ctx.save_for_backward(*saved) + ctx.is_varlen = cu_seqlens_q is not None + ctx.max_seqlen_q = max_seqlen_q + ctx.max_seqlen_k = max_seqlen_k + ctx.causal = causal + ctx.window_size = window_size + ctx.softmax_scale = softmax_scale if softmax_scale is not None else 1.0 / math.sqrt(q.shape[-1]) + ctx.softcap = softcap + ctx.deterministic = deterministic + return out_final + + @staticmethod + def backward(ctx, grad_output): + saved = ctx.saved_tensors + if ctx.is_varlen: + q, k, v, sink, raw_out, lse, cu_seqlens_q, cu_seqlens_k = saved + else: + q, k, v, sink, raw_out, lse = saved + cu_seqlens_q = cu_seqlens_k = None + + fa_kwargs = dict( + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=ctx.max_seqlen_q, + max_seqlen_k=ctx.max_seqlen_k, + softmax_scale=ctx.softmax_scale, + causal=ctx.causal, + window_size=ctx.window_size, + softcap=ctx.softcap, + deterministic=ctx.deterministic, + ) + scale = ctx.softmax_scale + + sink_f = sink.float().view(*([1] * (lse.ndim - 2)), -1, 1) + m = torch.sigmoid(lse - sink_f) # [..., Hq, S] + m_seq_head = m.transpose(-2, -1) # [..., S, Hq] + m_for_out = m_seq_head.unsqueeze(-1) # [..., S, Hq, 1] + + # fp32 accumulation is critical for dsink precision + g_r = (grad_output.float() * raw_out.float()).sum(dim=-1) # [..., S, Hq] + g_ell = g_r * m_seq_head * (1.0 - m_seq_head) # fp32 + + # --- path 1: main grad via FA bwd with dout' = dO * m --- + dout_main = (grad_output * m_for_out.to(grad_output.dtype)).contiguous() + dq_main = torch.empty_like(q) + dk_main = torch.empty_like(k) + dv = torch.empty_like(v) + _fa3_backward(dout_main, q, k, v, raw_out, lse, dq_main, dk_main, dv, **fa_kwargs) + + # --- path 2: dsink = -sum over all non-head axes of g_ell --- + dsink = -g_ell.flatten(0, -2).sum(dim=0).to(sink.dtype) # [Hq] + + # --- path 3: dq_extra = scale * g_ell * attention(Q, K, K) --- + mu_k = _fa3_forward(q, k, k, return_attn_probs=False, **fa_kwargs) + dq_extra = (scale * g_ell.unsqueeze(-1) * mu_k.float()).to(q.dtype) + + # --- path 4: dk_extra via FA bwd dv slot, dout' = g_ell * Q --- + x = (g_ell.unsqueeze(-1).to(q.dtype) * q).contiguous() + dq_dummy = torch.empty_like(q) + dk_dummy = torch.empty_like(k) + dk_extra = torch.empty_like(k) + _fa3_backward(x, q, k, k, raw_out, lse, dq_dummy, dk_dummy, dk_extra, **fa_kwargs) + dk_extra = scale * dk_extra + + dq = dq_main + dq_extra + dk = dk_main + dk_extra + # Nones correspond to: cu_q, cu_k, max_q, max_k, causal, ws, scale, sc, det + return dq, dk, dv, dsink, None, None, None, None, None, None, None, None, None + + +def flash_attn_with_sink( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + sink: torch.Tensor, + causal: bool = True, + window_size: Tuple[int, int] = (-1, -1), + softmax_scale: Optional[float] = None, + softcap: float = 0.0, + deterministic: bool = False, +) -> torch.Tensor: + """Batched FA3 attention with a per-head learned sink. + + Args: + q: ``[B, Sq, Hq, D]`` + k, v: ``[B, Sk, Hkv, D]`` (GQA supported) + sink: ``[Hq]`` fp32 learned per-head logit + """ + return FlashAttnWithSinkFA3.apply( + q, + k, + v, + sink, + None, + None, + None, + None, + causal, + window_size, + softmax_scale, + softcap, + deterministic, + ) + + +def flash_attn_varlen_with_sink( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + sink: torch.Tensor, + cu_seqlens_q: torch.Tensor, + cu_seqlens_k: torch.Tensor, + max_seqlen_q: int, + max_seqlen_k: int, + causal: bool = True, + window_size: Tuple[int, int] = (-1, -1), + softmax_scale: Optional[float] = None, + softcap: float = 0.0, + deterministic: bool = False, +) -> torch.Tensor: + """Varlen / packed-sequence FA3 attention with a per-head learned sink. + + Args: + q: ``[total_tokens, Hq, D]`` + k, v: ``[total_tokens, Hkv, D]`` + sink: ``[Hq]`` fp32 learned per-head logit + cu_seqlens_q, cu_seqlens_k: int32 prefix sums, shape ``[batch+1]`` + """ + return FlashAttnWithSinkFA3.apply( + q, + k, + v, + sink, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + causal, + window_size, + softmax_scale, + softcap, + deterministic, + ) diff --git a/src/xorl/models/transformers/gpt_oss/modeling_gpt_oss.py b/src/xorl/models/transformers/gpt_oss/modeling_gpt_oss.py new file mode 100644 index 00000000..97182335 --- /dev/null +++ b/src/xorl/models/transformers/gpt_oss/modeling_gpt_oss.py @@ -0,0 +1,858 @@ +"""GPT-OSS model implementation for the xorl framework. + +GPT-OSS is a Mixture-of-Experts transformer with: +- Standard RMSNorm (not zero-centered) +- Custom YaRN-style RoPE with NTK-by-parts scaling +- GQA with learned attention sinks (per-head softmax bias) +- Alternating sliding-window / full attention layers (even / odd) +- SwiGLU activation with interleaved layout, clamping, and alpha=1.702 +- Expert biases on both gate_up and down projections +- Router gate with bias +""" + +import math +from functools import partial +from typing import Callable, Optional, Tuple, Unpack + +import torch +import torch.nn.functional as F +from torch import nn + +from xorl.distributed.moe.deepep import sync_pending_combine +from xorl.distributed.parallel_state import get_parallel_state +from xorl.distributed.sequence_parallel.strategy import get_cp_strategy +from xorl.models.base import XorlPreTrainedModel +from xorl.models.layers.attention import AttentionKwargs, update_causal_mask +from xorl.models.layers.moe import MoEBlock +from xorl.models.layers.normalization import RMSNorm +from xorl.models.outputs import MoeCausalLMOutput, MoeModelOutput +from xorl.models.transformers.gpt_oss import parallelize +from xorl.models.transformers.gpt_oss.checkpoint_handler import GptOssCheckpointHandler +from xorl.models.transformers.gpt_oss.configuration_gpt_oss import GptOssConfig +from xorl.utils import logging + + +logger = logging.get_logger(__name__) + + +_sinks_grad_hook_logged = False + + +def _scale_sinks_grad_for_ulysses(grad: torch.Tensor) -> torch.Tensor: + global _sinks_grad_hook_logged + ps = get_parallel_state() + if ps.ulysses_enabled: + if not _sinks_grad_hook_logged: + logger.info(f"sinks grad hook active: scaling by ulysses_size={ps.ulysses_size}") + _sinks_grad_hook_logged = True + return grad * ps.ulysses_size + return grad + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _adapt_gpt_oss_config(config): + """Convert an HF config to GptOssConfig if needed.""" + if isinstance(config, GptOssConfig): + return config + return GptOssConfig.from_hf_config(config) + + +def _apply_rotary_emb(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: + """Apply rotary embeddings using half-dim split (not interleaved).""" + x1, x2 = x.chunk(2, dim=-1) + o1 = x1 * cos - x2 * sin + o2 = x2 * cos + x1 * sin + return torch.cat((o1, o2), dim=-1) + + +def gpt_oss_apply_rotary_pos_emb( + query: torch.Tensor, + key: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Apply rotary position embeddings to query and key. + + Args: + query: ``[batch, seq, num_heads, head_dim]`` + key: ``[batch, seq, num_kv_heads, head_dim]`` + cos: ``[batch, seq, head_dim // 2]`` + sin: ``[batch, seq, head_dim // 2]`` + """ + # Unsqueeze for head dimension: [batch, seq, 1, head_dim // 2] + cos = cos.unsqueeze(2).to(query.dtype) + sin = sin.unsqueeze(2).to(query.dtype) + return _apply_rotary_emb(query, cos, sin), _apply_rotary_emb(key, cos, sin) + + +def _repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """Expand key/value heads for grouped-query attention.""" + if n_rep == 1: + return hidden_states + batch, seq, num_kv_heads, head_dim = hidden_states.shape + hidden_states = hidden_states[:, :, :, None, :].expand(batch, seq, num_kv_heads, n_rep, head_dim) + return hidden_states.reshape(batch, seq, num_kv_heads * n_rep, head_dim) + + +# --------------------------------------------------------------------------- +# Rotary Embedding +# --------------------------------------------------------------------------- + + +class GptOssRotaryEmbedding(nn.Module): + """Custom YaRN-style RoPE with NTK-by-parts for GPT-OSS. + + See YaRN paper: https://arxiv.org/abs/2309.00071 + """ + + def __init__(self, config: GptOssConfig, device=None): + super().__init__() + self.head_dim = config.head_dim + self.base = config.rope_theta + self.initial_context_length = config.initial_context_length + self.scaling_factor = config.rope_scaling_factor + self.ntk_alpha = config.rope_ntk_alpha + self.ntk_beta = config.rope_ntk_beta + + def _compute_inv_freq_and_concentration(self, device: torch.device): + freq = self.base ** (torch.arange(0, self.head_dim, 2, dtype=torch.float32, device=device) / self.head_dim) + if self.scaling_factor > 1.0: + concentration = 0.1 * math.log(self.scaling_factor) + 1.0 + + d_half = self.head_dim / 2 + low = d_half * math.log(self.initial_context_length / (self.ntk_beta * 2 * math.pi)) / math.log(self.base) + high = d_half * math.log(self.initial_context_length / (self.ntk_alpha * 2 * math.pi)) / math.log(self.base) + + interpolation = 1.0 / (self.scaling_factor * freq) + extrapolation = 1.0 / freq + + ramp = (torch.arange(d_half, dtype=torch.float32, device=device) - low) / (high - low) + mask = 1 - ramp.clamp(0, 1) + + inv_freq = interpolation * (1 - mask) + extrapolation * mask + else: + concentration = 1.0 + inv_freq = 1.0 / freq + + return inv_freq, concentration + + @torch.no_grad() + def forward(self, x: torch.Tensor, position_ids: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + """Compute ``(cos, sin)`` position embeddings. + + Args: + x: Hidden states (used for device only). + position_ids: ``[batch, seq_len]``. + + Returns: + ``(cos, sin)`` each of shape ``[batch, seq_len, head_dim // 2]``. + """ + inv_freq, concentration = self._compute_inv_freq_and_concentration(x.device) + freqs = torch.einsum("bi,j->bij", position_ids.float(), inv_freq) + cos = freqs.cos() * concentration + sin = freqs.sin() * concentration + return cos, sin + + +# --------------------------------------------------------------------------- +# Attention +# --------------------------------------------------------------------------- + + +class GptOssAttention(nn.Module): + """GPT-OSS multi-head attention with learned sinks and alternating SWA. + + Attention sinks are per-head scalar biases that compete with real tokens + in the softmax, allowing the model to "waste" attention weight rather than + attending to specific tokens. Sinks are applied in both the eager path + (via explicit concat-then-softmax) and the ``flash_attention_3`` backend + (via a custom autograd wrapper that post-multiplies by + ``sigmoid(lse - sink)``; see ``flash_sink_attention.py``). Other backends + raise ``NotImplementedError`` — silently dropping sinks would corrupt + GPT-OSS semantics. + """ + + def __init__(self, config: GptOssConfig, layer_idx: int): + super().__init__() + self.config = config + self.layer_idx = layer_idx + self.head_dim = config.head_dim + self.num_attention_heads = config.num_attention_heads + self.num_key_value_heads = config.num_key_value_heads + self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads + self.scaling = self.head_dim**-0.5 + self.attention_dropout = config.attention_dropout + self.is_causal = True + + # Alternating sliding window: even layers → SWA, odd layers → full + self.sliding_window = config.sliding_window if layer_idx % 2 == 0 else None + + # Learned attention sinks (one per query head) + self.sinks = nn.Parameter(torch.empty(config.num_attention_heads)) + # Ulysses SP shards heads: only ulysses_rank=r writes grad to slice r, + # but FSDP averages across the full dp×ulysses group. Scale by + # ulysses_size so the post-reduce result is the correct dp-only average. + self.sinks.register_hook(_scale_sinks_grad_for_ulysses) + + # Fused QKV projection (with bias) + qkv_dim = config.head_dim * (config.num_attention_heads + 2 * config.num_key_value_heads) + self.qkv_proj = nn.Linear(config.hidden_size, qkv_dim, bias=config.attention_bias) + + # Output projection (with bias) + self.o_proj = nn.Linear( + config.num_attention_heads * config.head_dim, + config.hidden_size, + bias=config.attention_bias, + ) + + def _local_sinks(self) -> torch.Tensor: + # Ulysses SP shards heads across ranks via all-to-all; the sinks + # parameter is replicated, so slice it to match this rank's head shard. + ps = get_parallel_state() + if not ps.ulysses_enabled: + return self.sinks + sp = ps.ulysses_size + heads_per_rank = self.sinks.size(0) // sp + rank = ps.ulysses_rank + return self.sinks[rank * heads_per_rank : (rank + 1) * heads_per_rank] + + # -- TP helpers ---------------------------------------------------------- + + def unfuse_for_tp(self): + """Split fused QKV projection into separate q/k/v for tensor parallelism.""" + device = self.qkv_proj.weight.device + dtype = self.qkv_proj.weight.dtype + has_bias = self.qkv_proj.bias is not None + q_dim = self.num_attention_heads * self.head_dim + kv_dim = self.num_key_value_heads * self.head_dim + + self.q_proj = nn.Linear(self.config.hidden_size, q_dim, bias=has_bias, device=device, dtype=dtype) + self.k_proj = nn.Linear(self.config.hidden_size, kv_dim, bias=has_bias, device=device, dtype=dtype) + self.v_proj = nn.Linear(self.config.hidden_size, kv_dim, bias=has_bias, device=device, dtype=dtype) + + self.q_proj.weight.data.copy_(self.qkv_proj.weight.data[:q_dim]) + self.k_proj.weight.data.copy_(self.qkv_proj.weight.data[q_dim : q_dim + kv_dim]) + self.v_proj.weight.data.copy_(self.qkv_proj.weight.data[q_dim + kv_dim :]) + if has_bias: + self.q_proj.bias.data.copy_(self.qkv_proj.bias.data[:q_dim]) + self.k_proj.bias.data.copy_(self.qkv_proj.bias.data[q_dim : q_dim + kv_dim]) + self.v_proj.bias.data.copy_(self.qkv_proj.bias.data[q_dim + kv_dim :]) + + del self.qkv_proj + + # -- Attention strategy hooks ------------------------------------------- + + def _project_qkv( + self, + hidden_states: torch.Tensor, + position_embeddings: Tuple[torch.Tensor, torch.Tensor], + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + input_shape = hidden_states.shape[:-1] + hidden_shape = (*input_shape, -1, self.head_dim) + + if hasattr(self, "qkv_proj"): + qkv = self.qkv_proj(hidden_states) + q_dim = self.num_attention_heads * self.head_dim + kv_dim = self.num_key_value_heads * self.head_dim + query_states = qkv[..., :q_dim].view(hidden_shape) + key_states = qkv[..., q_dim : q_dim + kv_dim].view(hidden_shape) + value_states = qkv[..., q_dim + kv_dim :].view(hidden_shape) + else: + query_states = self.q_proj(hidden_states).view(hidden_shape) + key_states = self.k_proj(hidden_states).view(hidden_shape) + value_states = self.v_proj(hidden_states).view(hidden_shape) + + cos, sin = position_embeddings + query_states, key_states = gpt_oss_apply_rotary_pos_emb(query_states, key_states, cos, sin) + return query_states, key_states, value_states + + def _project_output(self, attn_output: torch.Tensor) -> torch.Tensor: + attn_output = attn_output.reshape(*attn_output.shape[:-2], -1).contiguous() + return self.o_proj(attn_output) + + def _get_attention_fn(self) -> Callable: + impl = getattr(self.config, "_attn_implementation", "eager") + if impl == "eager": + return self._attention_with_sinks + if impl == "flash_attention_3": + return self._flash_attention_3_with_sinks + if impl == "flash_attention_4": + raise NotImplementedError( + "GPT-OSS attention sinks are not wired through the FA4 (CUTE) backend. Use flash_attention_3 or eager." + ) + if impl in ATTENTION_FUNCTIONS: + raise NotImplementedError( + f"GPT-OSS attention sinks are not wired through the {impl!r} backend. Use flash_attention_3 or eager." + ) + return ATTENTION_FUNCTIONS.get(impl, self._attention_with_sinks) + + def _flash_attention_3_with_sinks( + self, + module: nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: Optional[torch.Tensor], + dropout: float = 0.0, + scaling: Optional[float] = None, + sliding_window: Optional[int] = None, + softcap: Optional[float] = None, + **kwargs, + ) -> Tuple[torch.Tensor, None]: + """FA3 forward with learned sinks applied via sigmoid(lse - sink). + + The installed FA3 build does not expose a native ``sinks=`` kwarg, so + sinks are fused by a custom autograd wrapper (FA3 fwd + manual bwd). + Supports both batched 4D and packed-varlen 3D paths. + """ + from xorl.models.transformers.gpt_oss.flash_sink_attention import ( + flash_attn_varlen_with_sink, + flash_attn_with_sink, + ) + + causal = getattr(module, "is_causal", True) + if sliding_window is not None: + window_size = (sliding_window, 0 if causal else sliding_window) + else: + window_size = (-1, -1) + + cu_seq_lens_q = kwargs.get("cu_seq_lens_q", None) + cu_seq_lens_k = kwargs.get("cu_seq_lens_k", None) + + if cu_seq_lens_q is not None and cu_seq_lens_k is not None: + # Packed / varlen path — flatten batch dim (packing collator uses B=1). + cu_seq_lens_q = cu_seq_lens_q.to(torch.int32) + cu_seq_lens_k = cu_seq_lens_k.to(torch.int32) + max_length_q = kwargs.get("max_length_q") + max_length_k = kwargs.get("max_length_k") + + def _flatten(x): + return x.squeeze(0) if x.size(0) == 1 else x.reshape(-1, x.size(-2), x.size(-1)) + + q_varlen = _flatten(query) + k_varlen = _flatten(key) + v_varlen = _flatten(value) + + out = flash_attn_varlen_with_sink( + q_varlen, + k_varlen, + v_varlen, + self._local_sinks(), + cu_seqlens_q=cu_seq_lens_q, + cu_seqlens_k=cu_seq_lens_k, + max_seqlen_q=max_length_q, + max_seqlen_k=max_length_k, + causal=causal, + window_size=window_size, + softmax_scale=scaling, + softcap=softcap if softcap is not None else 0.0, + ) + out = out.unsqueeze(0) + else: + out = flash_attn_with_sink( + query, + key, + value, + self._local_sinks(), + causal=causal, + window_size=window_size, + softmax_scale=scaling, + softcap=softcap if softcap is not None else 0.0, + ) + return out, None + + def _attention_kwargs(self) -> dict: + return dict( + dropout=0.0 if not self.training else self.attention_dropout, + scaling=self.scaling, + sliding_window=self.sliding_window, + ) + + # -- Eager attention with sinks ----------------------------------------- + + def _attention_with_sinks( + self, + module: nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: Optional[torch.Tensor], + scaling: float, + dropout: float = 0.0, + sliding_window: Optional[int] = None, + **kwargs, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + """Eager attention that injects learned sink logits into softmax. + + Args: + query: ``[batch, seq, num_heads, head_dim]`` + key: ``[batch, seq, num_kv_heads, head_dim]`` + value: ``[batch, seq, num_kv_heads, head_dim]`` + attention_mask: ``[batch, 1, seq, seq]`` or *None* + sliding_window: Per-layer sliding window size or *None*. + """ + key = _repeat_kv(key, self.num_key_value_groups) + value = _repeat_kv(value, self.num_key_value_groups) + + # -> [batch, heads, seq, head_dim] + query = query.transpose(1, 2) + key = key.transpose(1, 2) + value = value.transpose(1, 2) + + attn_weights = torch.matmul(query, key.transpose(-2, -1)) * scaling + + # Causal mask — trim to actual key length (framework mask may + # include an extra column for KV-cache bookkeeping). + if attention_mask is not None: + causal_mask = attention_mask + mk = causal_mask.shape[-1] + k_len = attn_weights.shape[-1] + if mk != k_len: + causal_mask = causal_mask[..., :k_len] + attn_weights = attn_weights + causal_mask + + # Sliding window mask (lower triangle beyond window) + if sliding_window is not None and sliding_window > 0: + seq_len = attn_weights.shape[-1] + sw_mask = torch.tril( + attn_weights.new_ones(seq_len, seq_len, dtype=torch.bool), + diagonal=-sliding_window, + ) + attn_weights = attn_weights.masked_fill(sw_mask, float("-inf")) + + # Attention sinks: [num_heads] -> [1, num_heads, 1, 1] -> [B, H, S, 1] + sinks = self._local_sinks().to(attn_weights.dtype) + sink_logits = sinks.view(1, -1, 1, 1).expand(attn_weights.shape[0], -1, attn_weights.shape[2], 1) + attn_weights = torch.cat([attn_weights, sink_logits], dim=-1) + + # Subtract max for numerical stability (matches HF implementation) + attn_weights = attn_weights - attn_weights.max(dim=-1, keepdim=True).values + # Softmax over (seq + 1) then drop the sink column + attn_weights = F.softmax(attn_weights, dim=-1, dtype=attn_weights.dtype) + attn_weights = attn_weights[..., :-1] + + attn_weights = F.dropout(attn_weights, p=dropout, training=module.training) + attn_output = torch.matmul(attn_weights, value) + attn_output = attn_output.transpose(1, 2).contiguous() + return attn_output, None + + # -- Forward ------------------------------------------------------------- + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: Tuple[torch.Tensor, torch.Tensor], + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values=None, + **kwargs: Unpack[AttentionKwargs], + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + 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.project_output(self, attn_output) + return attn_output, None + + +# --------------------------------------------------------------------------- +# MoE Block +# --------------------------------------------------------------------------- + + +class GptOssMoEBlock(MoEBlock): + """GPT-OSS MoE with expert biases and clamped SwiGLU activation. + + The checkpoint handler deinterleaves the original interleaved gate/up + layout into the standard concatenated ``[gate | up]`` format used by xorl + MoE backends. Per-expert biases (``gate_up_bias``, ``down_bias``) and the + ``"clamped_swiglu"`` activation kind are threaded through the shared + ``hidden_act`` dispatch (``SUPPORTED_HIDDEN_ACTS``). This supports both + single-GPU and Expert Parallel (EP) execution. + """ + + _SUPPORTED_MOE_IMPLEMENTATIONS = {"eager", "native"} + + def __init__(self, config: GptOssConfig, moe_implementation="triton"): + if moe_implementation not in self._SUPPORTED_MOE_IMPLEMENTATIONS: + raise NotImplementedError( + f"GPT-OSS requires per-expert biases (gate_up_bias, down_bias) and a " + f"clamped SwiGLU activation, which the {moe_implementation!r} MoE backend " + f"does not currently thread through. Supported backends: " + f"{sorted(self._SUPPORTED_MOE_IMPLEMENTATIONS)}. " + f"Running with {moe_implementation!r} would silently drop the biases and " + f"use a plain SiLU SwiGLU, producing incorrect outputs." + ) + super().__init__( + hidden_size=config.hidden_size, + num_experts=config.num_experts, + top_k=config.num_experts_per_tok, + intermediate_size=config.moe_intermediate_size, + 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) + + # GPT-OSS uses a clamped SwiGLU rather than the default SiLU SwiGLU. + # The activation is threaded through ``hidden_act`` so that all MoE + # backends that support it (currently ``eager`` and ``native``) dispatch + # on the string. Triton/quack raise NotImplementedError at entry. + self.experts.hidden_act = "clamped_swiglu" + + # GPT-OSS router gate has bias (base MoEBlock creates bias=False) + self.gate = nn.Linear(config.hidden_size, config.num_experts, bias=True) + + # Expert biases — registered on ``self.experts`` so that the + # checkpoint handler finds them at + # ``model.layers.*.mlp.experts.gate_up_bias`` / ``down_bias``. + self.experts.gate_up_bias = nn.Parameter(torch.zeros(config.num_experts, 2 * config.moe_intermediate_size)) + self.experts.down_bias = nn.Parameter(torch.zeros(config.num_experts, config.hidden_size)) + + def forward(self, hidden_states: torch.Tensor): + batch_size, sequence_length, hidden_dim = hidden_states.shape + hidden_states_flat = hidden_states.view(-1, hidden_dim) + + # Routing — topk on raw logits, then softmax on top-k values + # (matches the original OSS and HF implementations). + router_logits = self.gate(hidden_states_flat) + top_values, selected_experts = torch.topk(router_logits, k=self.top_k, dim=-1, sorted=True) + routing_weights = F.softmax(top_values, dim=1, dtype=top_values.dtype) + + if not self.train_router: + routing_weights = routing_weights.detach() + + if self.moe_implementation == "eager": + final = self._eager_forward(hidden_states_flat, routing_weights, selected_experts) + else: + final = self.experts(hidden_states_flat, routing_weights, selected_experts) + final = final.view(batch_size, sequence_length, hidden_dim) + return final, router_logits + + +GPT_OSS_MOE_CLASSES = { + "eager": partial(GptOssMoEBlock, moe_implementation="eager"), + "native": partial(GptOssMoEBlock, moe_implementation="native"), +} + + +# --------------------------------------------------------------------------- +# Decoder Layer +# --------------------------------------------------------------------------- + + +class GptOssDecoderLayer(nn.Module): + def __init__(self, config: GptOssConfig, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + self.self_attn = GptOssAttention(config, layer_idx) + moe_implementation = getattr(config, "_moe_implementation", "triton") + self.mlp = GPT_OSS_MOE_CLASSES[moe_implementation](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, + 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, ...]: + # Pre-norm → attention + 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_ids=position_ids, + past_key_values=past_key_values, + position_embeddings=position_embeddings, + **kwargs, + ) + + # Pre-norm → MLP (fused residual add via prenorm) + hidden_states, residual = self.post_attention_layernorm(hidden_states, residual=residual, prenorm=True) + hidden_states = self.mlp(hidden_states) + if isinstance(hidden_states, tuple): + hidden_states, router_logits = hidden_states + else: + router_logits = None + sync_pending_combine() + hidden_states = residual + hidden_states + + outputs = (hidden_states,) + if output_attentions: + outputs += (self_attn_weights,) + if output_router_logits: + outputs += (router_logits,) + return outputs + + +# --------------------------------------------------------------------------- +# Pre-trained model base +# --------------------------------------------------------------------------- + + +class GptOssPreTrainedModel(XorlPreTrainedModel): + config_class = GptOssConfig + base_model_prefix = "model" + _no_split_modules = ["GptOssDecoderLayer"] + + def _init_weights(self, module): + std = getattr(self.config, "initializer_range", 0.02) + 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) + + def get_parallel_plan(self): + return parallelize.get_ep_plan() + + def get_checkpoint_handler(self, **kwargs): + checkpoint_keys = kwargs.get("checkpoint_keys", set()) + ep_rank = kwargs.get("ep_rank", 0) + ep_size = kwargs.get("ep_size", 1) + is_broadcast = kwargs.get("is_broadcast", False) + if is_broadcast: + ep_rank, ep_size = 0, 1 + head_dim = getattr(self.config, "head_dim", self.config.hidden_size // self.config.num_attention_heads) + return GptOssCheckpointHandler( + num_experts=self.config.num_experts, + num_attention_heads=self.config.num_attention_heads, + num_key_value_heads=self.config.num_key_value_heads, + head_dim=head_dim, + ep_rank=ep_rank, + ep_size=ep_size, + checkpoint_keys=checkpoint_keys or None, + skip_qkv_merge=getattr(self, "_unfused_for_tp", False), + ) + + +# --------------------------------------------------------------------------- +# Model (backbone) +# --------------------------------------------------------------------------- + + +class GptOssModel(GptOssPreTrainedModel): + def __init__(self, config: GptOssConfig): + config = _adapt_gpt_oss_config(config) + super().__init__(config) + self.padding_idx = getattr(config, "pad_token_id", None) + self.vocab_size = config.vocab_size + self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + self.layers = nn.ModuleList( + [GptOssDecoderLayer(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 = GptOssRotaryEmbedding(config=config) + self.gradient_checkpointing = False + 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: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + use_cache: bool | None = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + output_router_logits: Optional[bool] = None, + **kwargs: Unpack[AttentionKwargs], + ) -> MoeModelOutput: + output_attentions = output_attentions if output_attentions is not None else False + output_router_logits = ( + output_router_logits + if output_router_logits is not None + else getattr(self.config, "output_router_logits", False) + ) + + 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 use_cache is None: + use_cache = False + + cache_position = torch.arange(hidden_states.shape[1], device=hidden_states.device) + causal_mask = update_causal_mask( + getattr(self.config, "_attn_implementation", "eager"), + attention_mask, + hidden_states, + cache_position, + sliding_window=None, # Per-layer SWA handled in the attention module + is_training=self.training, + output_attentions=output_attentions, + ) + + ps = get_parallel_state() + position_embeddings = self.rotary_emb(hidden_states, position_ids) + position_embeddings = get_cp_strategy().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 + all_router_logits = () if output_router_logits else None + + for decoder_layer in self.layers: + if decoder_layer is None: + continue + if self.gradient_checkpointing and self.training: + layer_outputs = self._gradient_checkpointing_func( + decoder_layer.__call__, + hidden_states, + causal_mask, + position_ids, + None, + use_cache, + output_attentions, + output_router_logits, + position_embeddings, + **kwargs, + ) + else: + layer_outputs = decoder_layer( + hidden_states, + attention_mask=causal_mask, + position_ids=position_ids, + past_key_values=None, + use_cache=use_cache, + output_attentions=output_attentions, + output_router_logits=output_router_logits, + position_embeddings=position_embeddings, + **kwargs, + ) + hidden_states = layer_outputs[0] + if output_attentions: + all_self_attns += (layer_outputs[1],) + 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, + ) + + +# --------------------------------------------------------------------------- +# Causal LM +# --------------------------------------------------------------------------- + + +class GptOssForCausalLM(GptOssPreTrainedModel): + _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_gpt_oss_config(config) + super().__init__(config) + self.model = GptOssModel(config) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.num_experts = config.num_experts + self.num_experts_per_tok = config.num_experts_per_tok + self.post_init() + + def unfuse_for_tp(self): + 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: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + **kwargs, + ) -> MoeCausalLMOutput: + output_router_logits = getattr(self.config, "output_router_logits", False) + outputs = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + output_router_logits=output_router_logits, + **kwargs, + ) + return MoeCausalLMOutput( + last_hidden_state=outputs.last_hidden_state, + router_logits=outputs.router_logits, + ) + + +ModelClass = [GptOssForCausalLM] + +__all__ = [ + "GptOssForCausalLM", + "GptOssModel", + "GptOssPreTrainedModel", + "GptOssMoEBlock", +] diff --git a/src/xorl/models/transformers/gpt_oss/parallelize.py b/src/xorl/models/transformers/gpt_oss/parallelize.py new file mode 100644 index 00000000..7d960bf6 --- /dev/null +++ b/src/xorl/models/transformers/gpt_oss/parallelize.py @@ -0,0 +1,52 @@ +"""Parallelization plan and utilities for GPT-OSS models.""" + +from torch.distributed._tensor import Shard + +from ....distributed.parallel_plan import ParallelPlan + + +# TP plan for the base model (GptOssModel). +# Only covers attention — all MLP layers are MoE and use EP instead. +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", +} + +# TP plan for top-level modules on the CausalLM wrapper. +MODEL_TP_PLAN = { + "lm_head": "colwise_rep", +} + + +def unfuse_for_tp(model): + """Unfuse fused QKV projections for tensor parallelism compatibility. + + GPT-OSS has only MoE layers (no dense MLP), so only attention is unfused. + """ + for layer in model.model.layers: + layer.self_attn.unfuse_for_tp() + model._unfused_for_tp = True + model.config.base_model_tp_plan = TP_PLAN + + +def get_ep_plan(): + """Get EP (expert parallelism) plan for GPT-OSS model.""" + ep_plan = { + # 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), + # Expert biases (stacked [num_experts, ...] format) + "model.layers.*.mlp.experts.gate_up_bias": Shard(0), + "model.layers.*.mlp.experts.down_bias": Shard(0), + # LoRA weights for experts + "model.layers.*.mlp.experts.gate_proj_lora_A": Shard(0), + "model.layers.*.mlp.experts.gate_proj_lora_B": Shard(0), + "model.layers.*.mlp.experts.up_proj_lora_A": Shard(0), + "model.layers.*.mlp.experts.up_proj_lora_B": Shard(0), + "model.layers.*.mlp.experts.down_proj_lora_A": Shard(0), + "model.layers.*.mlp.experts.down_proj_lora_B": Shard(0), + } + return ParallelPlan(ep_plan=ep_plan) diff --git a/src/xorl/models/transformers/olmo2/__init__.py b/src/xorl/models/transformers/olmo2/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/xorl/models/transformers/olmo2/checkpoint_handler.py b/src/xorl/models/transformers/olmo2/checkpoint_handler.py new file mode 100644 index 00000000..c07d8a50 --- /dev/null +++ b/src/xorl/models/transformers/olmo2/checkpoint_handler.py @@ -0,0 +1,170 @@ +"""Checkpoint handler for OLMo-2 dense 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 Olmo2CheckpointHandler(CheckpointHandler): + """Checkpoint handler for OLMo-2 dense 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 + + OLMo-2 layer norm names (``post_attention_layernorm``, + ``post_feedforward_layernorm``) already match the model's parameter + names, so no key remapping is needed for them. + """ + + 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/olmo2/configuration_olmo2.py b/src/xorl/models/transformers/olmo2/configuration_olmo2.py new file mode 100644 index 00000000..4250ce86 --- /dev/null +++ b/src/xorl/models/transformers/olmo2/configuration_olmo2.py @@ -0,0 +1,84 @@ +"""OLMo-2 model configuration.""" + +from transformers.configuration_utils import PretrainedConfig + +from ....utils import logging +from .parallelize import TP_PLAN + + +logger = logging.get_logger(__name__) + + +class Olmo2Config(PretrainedConfig): + r""" + Configuration class for the OLMo-2 dense model. + + OLMo-2 differs from Llama in two ways: + * Layer norms are applied **after** attention/MLP (post-norm), not before. + * QK normalization is applied across the full ``num_heads * head_dim`` + axis prior to reshape, rather than per-head. + """ + + model_type = "olmo2" + + 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=100352, + hidden_size=4096, + intermediate_size=11008, + num_hidden_layers=32, + num_attention_heads=32, + num_key_value_heads=32, + hidden_act="silu", + max_position_embeddings=4096, + initializer_range=0.02, + rms_norm_eps=1e-6, + use_cache=True, + tie_word_embeddings=False, + rope_theta=500000.0, + rope_scaling=None, + attention_bias=False, + attention_dropout=0.0, + pad_token_id=1, + bos_token_id=None, + eos_token_id=50279, + **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 = 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.rope_theta = rope_theta + self.rope_scaling = rope_scaling + self.attention_bias = attention_bias + self.attention_dropout = attention_dropout + + # 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__( + 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, + ) + + +__all__ = ["Olmo2Config"] diff --git a/src/xorl/models/transformers/olmo2/modeling_olmo2.py b/src/xorl/models/transformers/olmo2/modeling_olmo2.py new file mode 100644 index 00000000..2f39d74c --- /dev/null +++ b/src/xorl/models/transformers/olmo2/modeling_olmo2.py @@ -0,0 +1,424 @@ +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.layers.normalization import native_rms_norm +from xorl.models.layers.rope import apply_rotary_pos_emb +from xorl.models.module_utils import GradientCheckpointingLayer +from xorl.models.outputs import BaseModelOutput, CausalLMOutput +from xorl.models.transformers.olmo2 import parallelize +from xorl.models.transformers.olmo2.checkpoint_handler import Olmo2CheckpointHandler +from xorl.models.transformers.olmo2.configuration_olmo2 import Olmo2Config +from xorl.ops.fused_silu_and_mul import fused_silu_and_mul +from xorl.utils import logging + + +logger = logging.get_logger(__name__) + + +class Olmo2MLP(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 Olmo2QKRMSNorm(RMSNorm): + """Full-axis RMSNorm for OLMo-2 ``q_norm``/``k_norm`` under colwise TP. + + Without TP the parent ``forward`` runs unchanged. Under TP the OLMo-2 plan + applies ``LocalAxisRMSNormShard`` to this module's weight, sharding it on + dim 0 across the TP mesh. The colwise ``q_proj``/``k_proj`` produces a + plain local tensor whose last dim already equals the rank's weight slice + (``num_heads * head_dim / tp``). ``F.rms_norm`` has no DTensor sharding + rule, so dispatch through ``__torch_function__`` would mis-handle the + sharded weight; bypass it by running the fused op directly on the local + weight tensor. The rank-local RMS this computes matches HuggingFace's + ``Olmo2RMSNorm`` reference behavior under TP (a deliberate + local-vs-global RMS approximation, since the partition axis IS the norm + axis for this model). + """ + + def forward(self, hidden_states, residual=None, prenorm=False): + from torch.distributed.tensor import DTensor # noqa: PLC0415 + + weight = self.weight + if not isinstance(weight, DTensor): + return super().forward(hidden_states, residual, prenorm) + + residual_out = None + norm_input = hidden_states + if residual is not None: + residual_out = hidden_states + residual + norm_input = residual_out + + local_weight = weight.to_local() + out = native_rms_norm(norm_input, local_weight, self.variance_epsilon) + if residual_out is not None and prenorm: + return out, residual_out + return out + + +class Olmo2Attention(MultiHeadAttention): + """OLMo-2 attention. + + OLMo-2 normalizes Q and K across the full ``num_heads * head_dim`` axis + (not per-head as in Qwen3). The base ``MultiHeadAttention`` allocates + per-head q_norm/k_norm when ``use_qk_norm=True``; we set it to ``False`` + on the config seen by the base class and own the full-axis norms here. + """ + + def __init__(self, config, layer_idx: int): + # Disable base-class per-head q_norm/k_norm; we install full-axis norms. + config.use_qk_norm = False + super().__init__(config, layer_idx) + self.q_norm = Olmo2QKRMSNorm(config.num_attention_heads * self.head_dim, eps=config.rms_norm_eps) + self.k_norm = Olmo2QKRMSNorm(config.num_key_value_heads * self.head_dim, eps=config.rms_norm_eps) + + def _init_sliding_window(self, config): + return None + + def _project_qkv( + self, + hidden_states: torch.Tensor, + position_embeddings: Tuple[torch.Tensor, torch.Tensor], + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + input_shape = hidden_states.shape[:-1] + hidden_shape = (*input_shape, -1, self.head_dim) + + if hasattr(self, "qkv_proj"): + qkv = self.qkv_proj(hidden_states) + q, k, v = qkv.split([self.q_dim, self.kv_dim, self.kv_dim], dim=-1) + else: + q = self.q_proj(hidden_states) + k = self.k_proj(hidden_states) + v = self.v_proj(hidden_states) + + # Full-axis QK norm before reshape (OLMo-2 specific). + q = self.q_norm(q) + k = self.k_norm(k) + + q = q.view(hidden_shape) + k = k.view(hidden_shape) + v = v.view(hidden_shape) + + cos, sin = position_embeddings + q, k = apply_rotary_pos_emb(q, k, cos, sin) + + if getattr(self.config, "_attention_cast_bf16", False): + q = q.to(torch.bfloat16) + k = k.to(torch.bfloat16) + + return q, k, v + + +class Olmo2DecoderLayer(GradientCheckpointingLayer): + """OLMo-2 decoder layer with post-attention and post-feedforward norms. + + Unlike Llama's pre-norm, OLMo-2 normalizes after each sublayer and adds + the residual afterwards. There is no ``input_layernorm`` -- the only + norms are ``post_attention_layernorm`` and ``post_feedforward_layernorm``. + """ + + def __init__(self, config: Olmo2Config, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + self.self_attn = Olmo2Attention(config=config, layer_idx=layer_idx) + self.mlp = Olmo2MLP(config) + self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_feedforward_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 + + # Self attention -> post-norm -> residual add + hidden_states, self_attn_weights = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_embeddings=position_embeddings, + **kwargs, + ) + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = residual + hidden_states + + # MLP -> post-norm -> residual add + residual = hidden_states + hidden_states = self.mlp(hidden_states) + hidden_states = self.post_feedforward_layernorm(hidden_states) + hidden_states = residual + hidden_states + + outputs = (hidden_states,) + if output_attentions: + outputs += (self_attn_weights,) + + return outputs + + +class Olmo2PreTrainedModel(XorlPreTrainedModel): + config_class = Olmo2Config + base_model_prefix = "model" + _no_split_modules = ["Olmo2DecoderLayer"] + + 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 Olmo2CheckpointHandler( + 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 Olmo2Model(Olmo2PreTrainedModel): + """OLMo-2 transformer decoder.""" + + def __init__(self, config: Olmo2Config): + 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( + [Olmo2DecoderLayer(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 + ) + + 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, + ) + + 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 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],) + + 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 Olmo2ForCausalLM(Olmo2PreTrainedModel): + _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 = Olmo2Model(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 = Olmo2ForCausalLM + +__all__ = ["Olmo2ForCausalLM", "Olmo2Model", "Olmo2PreTrainedModel"] diff --git a/src/xorl/models/transformers/olmo2/parallelize.py b/src/xorl/models/transformers/olmo2/parallelize.py new file mode 100644 index 00000000..dfc34c98 --- /dev/null +++ b/src/xorl/models/transformers/olmo2/parallelize.py @@ -0,0 +1,80 @@ +"""Parallelization plan and utilities for OLMo-2 models. + +The plan keeps the residual stream **Replicate** across decoder layers +(rather than the torchtitan SP/Shard(1) pattern). This composes with +xorl's existing loss path: ``vocab_parallel_cross_entropy`` bypasses +``lm_head`` and matmuls ``hidden_states @ lm_head.weight.t()`` directly, +which expects a full ``[B, S, H]`` per rank — incompatible with a +sequence-sharded residual. + +The TP boundaries are therefore the standard colwise/rowwise pair: + + * Embedding outputs Replicate (default). + * q/k/v_proj and gate/up_proj are ``ColwiseParallel()`` (default + Replicate input, Shard(-1) output, ``use_local_output=True``) — the + block sees plain local tensors with ``hidden / tp`` per rank, so + rotary, flash-attention and the rest of the attention/MLP internals + run unchanged on local tensors. + * o_proj/down_proj are ``RowwiseParallel()`` (default Shard(-1) input, + Replicate output) — every block returns full hidden, so post-norms + and the residual addition see Replicate without any extra plumbing. + * lm_head is ``ColwiseParallel()`` (default Replicate input, Shard(-1) + output) for use with vocab-parallel cross-entropy. + +OLMo-2's full-axis ``q_norm``/``k_norm`` (over ``num_heads * head_dim``, +not per-head) doesn't compose with the stock styles: under colwise q/k_proj +the q/k tensors arrive with a sharded hidden axis, so a full-hidden weight +can't be applied directly. ``LocalAxisRMSNormShard`` (in ``tp_styles.py``) +shards the 1-D weight on dim 0 so each rank's slice matches its local q/k +slice; ``Olmo2QKRMSNorm.forward`` (in ``modeling_olmo2.py``) detects the +sharded weight and runs the fused op on locals — computing a local-axis +RMS that matches HuggingFace's ``Olmo2RMSNorm`` reference under TP. +""" + +from torch.distributed.tensor.parallel import ColwiseParallel, RowwiseParallel + +from xorl.models.transformers.olmo2.tp_styles import LocalAxisRMSNormShard + + +# Plan for the base model (Olmo2Model). Keys use wildcard patterns relative +# to the base model prefix, which the parallel applier prepends with ``model.``. +TP_PLAN = { + "embed_tokens": "embedding", + "layers.*.self_attn.q_proj": ColwiseParallel(), + "layers.*.self_attn.k_proj": ColwiseParallel(), + "layers.*.self_attn.v_proj": ColwiseParallel(), + # Full-axis QK norms: weight sharded on dim 0 so each rank's slice + # matches the local hidden slice from colwise q/k_proj. + "layers.*.self_attn.q_norm": LocalAxisRMSNormShard(), + "layers.*.self_attn.k_norm": LocalAxisRMSNormShard(), + "layers.*.self_attn.o_proj": RowwiseParallel(), + # post_attention_layernorm sees a Replicate input after the rowwise + # all-reduce in o_proj — no plan entry needed. + "layers.*.mlp.gate_proj": ColwiseParallel(), + "layers.*.mlp.up_proj": ColwiseParallel(), + "layers.*.mlp.down_proj": RowwiseParallel(), + # post_feedforward_layernorm and the final norm are also Replicate-fed. +} + +# Plan for top-level modules on the CausalLM wrapper. +MODEL_TP_PLAN = { + "lm_head": ColwiseParallel(), +} + + +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/olmo2/tp_styles.py b/src/xorl/models/transformers/olmo2/tp_styles.py new file mode 100644 index 00000000..91d59f0f --- /dev/null +++ b/src/xorl/models/transformers/olmo2/tp_styles.py @@ -0,0 +1,37 @@ +"""OLMo-2-specific tensor-parallel ``ParallelStyle``. + +OLMo-2 declares ``q_norm``/``k_norm`` over the full ``num_heads * head_dim`` +axis (rather than per-head + reshape-first like every other model in the +repo). Under colwise ``q_proj``/``k_proj``, the q/k tensors arrive with a +sharded hidden axis, so a full-hidden RMSNorm weight can't be applied +directly. ``LocalAxisRMSNormShard`` shards the 1-D RMSNorm weight along +dim 0 so each rank's slice matches its local q/k slice; the +DTensor-aware ``Olmo2QKRMSNorm.forward`` (in ``modeling_olmo2.py``) then +runs the fused op on locals and computes a local-axis RMS — matching +HuggingFace's ``Olmo2RMSNorm`` reference behavior. + +This style does not compose with the other models in this repo: per-head +QK norm reshape-first models stay on stock ``ColwiseParallel`` (their +norm weight is ``[head_dim]`` and the post-reshape activation has full +``head_dim`` per rank). The custom style is therefore scoped to +``olmo2/`` and not exported as generic TP infrastructure. +""" + +from torch import nn +from torch.distributed.tensor import Shard, distribute_tensor +from torch.distributed.tensor.parallel import ParallelStyle + + +class LocalAxisRMSNormShard(ParallelStyle): + """Shard a 1-D RMSNorm weight along dim 0 with no input/output redistribute. + + See module docstring for the OLMo-2 use case. The companion forward + is ``Olmo2QKRMSNorm.forward`` which detects the Shard(0) DTensor weight + and runs ``F.rms_norm`` on local tensors. + """ + + def _apply(self, module: nn.Module, device_mesh) -> nn.Module: + for name, param in module.named_parameters(recurse=False): + sharded = nn.Parameter(distribute_tensor(param, device_mesh, [Shard(0)])) + module.register_parameter(name, sharded) + return module 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 2678284e..c519cf2d 100644 --- a/src/xorl/models/transformers/qwen3_5/configuration_qwen3_5.py +++ b/src/xorl/models/transformers/qwen3_5/configuration_qwen3_5.py @@ -21,7 +21,7 @@ def _cfg_to_dict(value): if isinstance(value, dict): return dict(value) if hasattr(value, "__dict__"): - return {k: v for k, v in vars(value).items()} + return dict(vars(value)) return value @@ -32,6 +32,29 @@ def _split_mrope_fields(value): return value_dict or None, mrope_interleaved, mrope_section +def _expand_layer_types(layer_types, num_hidden_layers, full_attention_interval): + if layer_types is not None: + layer_types = list(layer_types) + if len(layer_types) == num_hidden_layers: + return layer_types + if full_attention_interval: + return [ + "full_attention" if (layer_idx + 1) % full_attention_interval == 0 else "linear_attention" + for layer_idx in range(num_hidden_layers) + ] + raise ValueError( + f"`layer_types` must have {num_hidden_layers} entries when `full_attention_interval` is not set, " + f"got {len(layer_types)}." + ) + + if full_attention_interval: + return [ + "full_attention" if (layer_idx + 1) % full_attention_interval == 0 else "linear_attention" + for layer_idx in range(num_hidden_layers) + ] + return ["full_attention"] * num_hidden_layers + + class Qwen3_5Config(PretrainedConfig): model_type = "xorl_qwen3_5" @@ -119,15 +142,7 @@ def __init__( ) 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" - for layer_idx in range(num_hidden_layers) - ] - else: - layer_types = ["full_attention"] * num_hidden_layers - self.layer_types = list(layer_types) + self.layer_types = _expand_layer_types(layer_types, num_hidden_layers, full_attention_interval) if self._rope_scaling is not None and "type" in self._rope_scaling: self._rope_scaling["rope_type"] = self._rope_scaling["type"] 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 938db2d1..d7d69c29 100644 --- a/src/xorl/models/transformers/qwen3_5/modeling_qwen3_5.py +++ b/src/xorl/models/transformers/qwen3_5/modeling_qwen3_5.py @@ -169,13 +169,10 @@ def _project_qkv( value_states = self.v_proj(hidden_states).view(hidden_shape) cos, sin = position_embeddings - query_states, key_states = qwen3_5_apply_rotary_pos_emb( - query_states, - key_states, - cos, - sin, - interleaved=getattr(self.config, "mrope_interleaved", False), - ) + # `mrope_interleaved` controls T/H/W frequency mixing in cos/sin + # construction upstream, not the q/k rotation convention. q/k always + # use the standard half-rotate convention (HF/SGLang). + query_states, key_states = qwen3_5_apply_rotary_pos_emb(query_states, key_states, cos, sin) return query_states.contiguous(), key_states.contiguous(), value_states.contiguous() def _project_output(self, attn_output: torch.Tensor) -> torch.Tensor: 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 2112f9d4..a15d5fdb 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 @@ -34,7 +34,7 @@ def _cfg_to_dict(value): if isinstance(value, dict): return dict(value) if hasattr(value, "__dict__"): - return {k: v for k, v in vars(value).items()} + return dict(vars(value)) return value @@ -45,6 +45,29 @@ def _split_mrope_fields(value): return value_dict or None, mrope_interleaved, mrope_section +def _expand_layer_types(layer_types, num_hidden_layers, full_attention_interval): + if layer_types is not None: + layer_types = list(layer_types) + if len(layer_types) == num_hidden_layers: + return layer_types + if full_attention_interval: + return [ + "full_attention" if (layer_idx + 1) % full_attention_interval == 0 else "linear_attention" + for layer_idx in range(num_hidden_layers) + ] + raise ValueError( + f"`layer_types` must have {num_hidden_layers} entries when `full_attention_interval` is not set, " + f"got {len(layer_types)}." + ) + + if full_attention_interval: + return [ + "full_attention" if (layer_idx + 1) % full_attention_interval == 0 else "linear_attention" + for layer_idx in range(num_hidden_layers) + ] + return ["full_attention"] * num_hidden_layers + + class Qwen3_5MoeConfig(PretrainedConfig): model_type = "xorl_qwen3_5_moe" @@ -141,15 +164,7 @@ def __init__( ) 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" - for layer_idx in range(num_hidden_layers) - ] - else: - layer_types = ["full_attention"] * num_hidden_layers - self.layer_types = list(layer_types) + self.layer_types = _expand_layer_types(layer_types, num_hidden_layers, full_attention_interval) self.shared_expert_intermediate_size = intermediate_size if self._rope_scaling is not None and "type" in self._rope_scaling: 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 52a54b97..0e4ca35b 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 @@ -168,13 +168,10 @@ def _project_qkv( value_states = self.v_proj(hidden_states).view(hidden_shape) cos, sin = position_embeddings - query_states, key_states = qwen3_5_apply_rotary_pos_emb( - query_states, - key_states, - cos, - sin, - interleaved=getattr(self.config, "mrope_interleaved", False), - ) + # `mrope_interleaved` controls T/H/W frequency mixing in cos/sin + # construction upstream, not the q/k rotation convention. q/k always + # use the standard half-rotate convention (HF/SGLang). + query_states, key_states = qwen3_5_apply_rotary_pos_emb(query_states, key_states, cos, sin) return query_states, key_states, value_states def _project_output(self, attn_output: torch.Tensor) -> torch.Tensor: @@ -522,7 +519,7 @@ def forward( _use_outer_checkpoint = ( self.gradient_checkpointing and self.training - and getattr(self, "_gradient_checkpointing_method", "recompute_full_layer") == "recompute_full_layer" + and self._gradient_checkpointing_method == "recompute_full_layer" ) if _use_outer_checkpoint: layer_outputs = self._gradient_checkpointing_func( diff --git a/src/xorl/models/transformers/qwen3_5_shared.py b/src/xorl/models/transformers/qwen3_5_shared.py index d93a0860..f042c695 100644 --- a/src/xorl/models/transformers/qwen3_5_shared.py +++ b/src/xorl/models/transformers/qwen3_5_shared.py @@ -51,6 +51,20 @@ def qwen3_5_apply_rotary_pos_emb( sin: torch.Tensor, interleaved: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: + # `interleaved` describes the q/k feature-layout convention only. + # - `False` (default): standard half-rotate. Used by Qwen3.5/Qwen3.6 + # (HF/SGLang). Qwen's `mrope_interleaved` is about T/H/W frequency + # mixing in cos/sin construction and must NOT be plumbed in here. + # - `True`: pairwise rotation on adjacent (2i, 2i+1) features. Used by + # DeepSeek-V3 MLA decoupled RoPE when `rope_interleave=True`. + if interleaved: + # `RotaryEmbedding` emits cos/sin in halved layout + # [c0, c1, ..., c_{d/2-1}, c0, c1, ..., c_{d/2-1}]. The interleaved + # rotate_half rotates pair i at indices (2i, 2i+1), so cos/sin must be + # in interleaved layout [c0, c0, c1, c1, ...] for the math to line up. + half = cos.shape[-1] // 2 + cos = cos[..., :half].repeat_interleave(2, dim=-1) + sin = sin[..., :half].repeat_interleave(2, dim=-1) cos = cos.unsqueeze(2) sin = sin.unsqueeze(2) rotary_dim = cos.shape[-1] diff --git a/src/xorl/models/transformers/qwen3_moe/modeling_qwen3_moe.py b/src/xorl/models/transformers/qwen3_moe/modeling_qwen3_moe.py index 5e11651e..7c3c46fb 100644 --- a/src/xorl/models/transformers/qwen3_moe/modeling_qwen3_moe.py +++ b/src/xorl/models/transformers/qwen3_moe/modeling_qwen3_moe.py @@ -438,14 +438,14 @@ 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 + _grad_ckpt_method = self._gradient_checkpointing_method if _grad_ckpt_active else None for decoder_layer in self.layers: if decoder_layer is None: # PP: pruned layer continue - if _grad_ckpt_active and _grad_ckpt_method == "recompute_full_layer": + if _grad_ckpt_method == "recompute_full_layer": # Recompute entire layer in backward (including dispatch + combine) layer_outputs = self._gradient_checkpointing_func( decoder_layer.__call__, @@ -457,7 +457,7 @@ def forward( position_embeddings, **kwargs, ) - elif _grad_ckpt_active and _grad_ckpt_method == "recompute_before_dispatch": + elif _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( diff --git a/src/xorl/ops/group_gemm/kernel/lora_utils.py b/src/xorl/ops/group_gemm/kernel/lora_utils.py index b6c6a354..9482ab2e 100644 --- a/src/xorl/ops/group_gemm/kernel/lora_utils.py +++ b/src/xorl/ops/group_gemm/kernel/lora_utils.py @@ -37,16 +37,16 @@ def init_lora_weights_stacked( Returns: Tuple of (lora_A, lora_B) tensors: - - lora_A: Shape [num_experts, r, in_features] - - lora_B: Shape [num_experts, out_features, r] + - lora_A: Shape [num_experts, in_features, r] + - lora_B: Shape [num_experts, r, out_features] """ # lora_A: projects input to low-rank space - # Shape: [num_experts, r, in_features] - lora_A = torch.empty(num_experts, r, in_features, dtype=dtype, device=device) + # Shape: [num_experts, in_features, r] + lora_A = torch.empty(num_experts, in_features, r, dtype=dtype, device=device) # lora_B: projects from low-rank space to output - # Shape: [num_experts, out_features, r] - lora_B = torch.zeros(num_experts, out_features, r, dtype=dtype, device=device) + # Shape: [num_experts, r, out_features] + lora_B = torch.zeros(num_experts, r, out_features, dtype=dtype, device=device) # Initialize lora_A if init_method == "kaiming": @@ -88,20 +88,20 @@ def merge_lora_weights_stacked( ) -> torch.Tensor: """Merge LoRA weights into base weights. - Computes: W' = W + B @ A * scaling + Computes: W' = W + A @ B * scaling Args: - base_weight: Base weight tensor [num_experts, out_features, in_features] - lora_A: LoRA A tensor [num_experts, r, in_features] - lora_B: LoRA B tensor [num_experts, out_features, r] + base_weight: Base weight tensor [num_experts, in_features, out_features] + lora_A: LoRA A tensor [num_experts, in_features, r] + lora_B: LoRA B tensor [num_experts, r, out_features] scaling: LoRA scaling factor Returns: - Merged weight tensor [num_experts, out_features, in_features] + Merged weight tensor [num_experts, in_features, out_features] """ - # B @ A: [num_experts, out_features, r] @ [num_experts, r, in_features] - # = [num_experts, out_features, in_features] - delta_weight = torch.bmm(lora_B, lora_A) * scaling + # A @ B: [num_experts, in_features, r] @ [num_experts, r, out_features] + # = [num_experts, in_features, out_features] + delta_weight = torch.bmm(lora_A, lora_B) * scaling return base_weight + delta_weight @@ -113,18 +113,18 @@ def unmerge_lora_weights_stacked( ) -> torch.Tensor: """Unmerge LoRA weights from merged weights. - Computes: W = W' - B @ A * scaling + Computes: W = W' - A @ B * scaling Args: - merged_weight: Merged weight tensor [num_experts, out_features, in_features] - lora_A: LoRA A tensor [num_experts, r, in_features] - lora_B: LoRA B tensor [num_experts, out_features, r] + merged_weight: Merged weight tensor [num_experts, in_features, out_features] + lora_A: LoRA A tensor [num_experts, in_features, r] + lora_B: LoRA B tensor [num_experts, r, out_features] scaling: LoRA scaling factor Returns: - Base weight tensor [num_experts, out_features, in_features] + Base weight tensor [num_experts, in_features, out_features] """ - delta_weight = torch.bmm(lora_B, lora_A) * scaling + delta_weight = torch.bmm(lora_A, lora_B) * scaling return merged_weight - delta_weight @@ -135,14 +135,14 @@ def get_lora_delta_weight_stacked( ) -> torch.Tensor: """Compute the LoRA weight delta. - Computes: delta_W = B @ A * scaling + Computes: delta_W = A @ B * scaling Args: - lora_A: LoRA A tensor [num_experts, r, in_features] - lora_B: LoRA B tensor [num_experts, out_features, r] + lora_A: LoRA A tensor [num_experts, in_features, r] + lora_B: LoRA B tensor [num_experts, r, out_features] scaling: LoRA scaling factor Returns: - Delta weight tensor [num_experts, out_features, in_features] + Delta weight tensor [num_experts, in_features, out_features] """ - return torch.bmm(lora_B, lora_A) * scaling + return torch.bmm(lora_A, lora_B) * scaling diff --git a/src/xorl/ops/linear_attention/__init__.py b/src/xorl/ops/linear_attention/__init__.py index a10dcd3a..d3e5aa11 100644 --- a/src/xorl/ops/linear_attention/__init__.py +++ b/src/xorl/ops/linear_attention/__init__.py @@ -1,7 +1,8 @@ """Linear attention ops and layers used by Qwen3.5.""" +from fla.ops.gated_delta_rule import chunk_gated_delta_rule, fused_recurrent_gated_delta_rule + from .layers.gated_deltanet import GatedDeltaNet -from .ops.gated_delta_rule import chunk_gated_delta_rule, fused_recurrent_gated_delta_rule __all__ = [ diff --git a/src/xorl/ops/linear_attention/layers/gated_deltanet.py b/src/xorl/ops/linear_attention/layers/gated_deltanet.py index 2e832db2..018f21fc 100644 --- a/src/xorl/ops/linear_attention/layers/gated_deltanet.py +++ b/src/xorl/ops/linear_attention/layers/gated_deltanet.py @@ -9,14 +9,17 @@ import torch import torch.nn as nn from einops import rearrange, repeat -from torch.nn import functional as F - -from xorl.ops.linear_attention.layers.utils import get_unpad_data, index_first_axis, pad_input -from xorl.ops.linear_attention.modules import FusedRMSNormGated, RMSNorm, ShortConvolution -from xorl.ops.linear_attention.ops.gated_delta_rule import ( +from fla.modules import FusedRMSNormGated, RMSNorm +from fla.ops.gated_delta_rule import ( chunk_gated_delta_rule, fused_recurrent_gated_delta_rule, ) +from torch.nn import functional as F + +# `ShortConvolution` is kept local (not `fla.modules`): the xorl version carries the +# Ulysses conv-prefix halo exchange (`cp_context`) that upstream FLA does not implement. +from xorl.ops.linear_attention.layers.utils import get_unpad_data, index_first_axis, pad_input +from xorl.ops.linear_attention.modules import ShortConvolution class GatedDeltaNet(nn.Module): diff --git a/src/xorl/ops/linear_attention/modules/__init__.py b/src/xorl/ops/linear_attention/modules/__init__.py index 8dbea46c..2c2fdac1 100644 --- a/src/xorl/ops/linear_attention/modules/__init__.py +++ b/src/xorl/ops/linear_attention/modules/__init__.py @@ -1,10 +1,6 @@ -from .fused_norm_gate import FusedRMSNormGated -from .layernorm import RMSNorm from .short_conv import ShortConvolution __all__ = [ - "FusedRMSNormGated", - "RMSNorm", "ShortConvolution", ] diff --git a/src/xorl/ops/linear_attention/modules/fused_norm_gate.py b/src/xorl/ops/linear_attention/modules/fused_norm_gate.py deleted file mode 100644 index 1ba9920e..00000000 --- a/src/xorl/ops/linear_attention/modules/fused_norm_gate.py +++ /dev/null @@ -1,89 +0,0 @@ -from __future__ import annotations - -# 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 - - -def rms_norm_gated( - x: torch.Tensor, - g: torch.Tensor, - weight: torch.Tensor | None, - bias: torch.Tensor | None, - activation: str = "swish", - residual: torch.Tensor | None = None, - prenorm: bool = False, - residual_in_fp32: bool = False, - eps: float = 1e-6, -) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - del residual_in_fp32 - if activation not in {"swish", "silu", "sigmoid"}: - raise ValueError(f"Unsupported activation: {activation}") - - residual_out = x if residual is None else x + residual - norm_input = residual_out.float() - inv_rms = torch.rsqrt(norm_input.square().mean(dim=-1, keepdim=True) + eps) - y = norm_input * inv_rms - if weight is not None: - y = y * weight.float() - if bias is not None: - y = y + bias.float() - - gate = g.float() - if activation in {"swish", "silu"}: - y = y * gate * torch.sigmoid(gate) - else: - y = y * torch.sigmoid(gate) - - y = y.to(x.dtype) - return y if not prenorm else (y, residual_out) - - -class FusedRMSNormGated(nn.Module): - def __init__( - self, - hidden_size: int, - elementwise_affine: bool = True, - eps: float = 1e-5, - activation: str = "swish", - device: torch.device | None = None, - dtype: torch.dtype | None = None, - ) -> None: - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.hidden_size = hidden_size - self.elementwise_affine = elementwise_affine - self.eps = eps - self.activation = activation - - self.register_parameter("weight", None) - self.register_parameter("bias", None) - if elementwise_affine: - self.weight = nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - - self.reset_parameters() - - def reset_parameters(self) -> None: - if self.weight is not None: - nn.init.ones_(self.weight) - - def forward( - self, - x: torch.Tensor, - g: torch.Tensor, - residual: torch.Tensor | None = None, - prenorm: bool = False, - residual_in_fp32: bool = False, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - return rms_norm_gated( - x=x, - g=g, - weight=self.weight, - bias=self.bias, - activation=self.activation, - residual=residual, - prenorm=prenorm, - residual_in_fp32=residual_in_fp32, - eps=self.eps, - ) diff --git a/src/xorl/ops/linear_attention/modules/l2norm.py b/src/xorl/ops/linear_attention/modules/l2norm.py deleted file mode 100644 index ea5ed03d..00000000 --- a/src/xorl/ops/linear_attention/modules/l2norm.py +++ /dev/null @@ -1,266 +0,0 @@ -# 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 triton.language as tl - -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] - - -@triton.autotune( - configs=[triton.Config({}, num_warps=num_warps) for num_warps in NUM_WARPS_AUTOTUNE], - key=["D"], - **autotune_cache_kwargs, -) -@triton.jit -def l2norm_fwd_kernel1( - x, - y, - rstd, - eps, - D, - BD: tl.constexpr, -): - i_t = tl.program_id(0) - x += i_t * D - y += i_t * D - # Compute mean and variance - cols = tl.arange(0, BD) - mask = cols < D - - b_x = tl.load(x + cols, mask=mask, other=0.0).to(tl.float32) - b_rstd = 1 / tl.sqrt(tl.sum(b_x * b_x) + eps) - b_y = b_x * b_rstd - tl.store(y + cols, b_y, mask=mask) - tl.store(rstd + i_t, b_rstd) - - -@triton.autotune( - configs=[triton.Config({}, num_warps=num_warps) for num_warps in NUM_WARPS_AUTOTUNE], - key=["D"], - **autotune_cache_kwargs, -) -@triton.jit -def l2norm_bwd_kernel1( - y, - rstd, - dy, - dx, - eps, - D, - BD: tl.constexpr, -): - i_t = tl.program_id(0) - y += i_t * D - dx += i_t * D - dy += i_t * D - - cols = tl.arange(0, BD) - mask = cols < D - b_y = tl.load(y + cols, mask=mask, other=0.0).to(tl.float32) - b_rstd = tl.load(rstd + i_t).to(tl.float32) - b_dy = tl.load(dy + cols, mask=mask, other=0.0).to(tl.float32) - b_dx = b_dy * b_rstd - tl.sum(b_dy * b_y) * b_y * b_rstd - tl.store(dx + cols, b_dx, mask=mask) - - -@triton.autotune( - configs=[triton.Config({"BT": BT}, num_warps=num_warps) for num_warps in [1, 2, 4, 8, 16] for BT in BT_LIST], - key=["D", "NB"], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=["T"]) -def l2norm_fwd_kernel( - x, - y, - rstd, - eps, - T, - D: tl.constexpr, - BD: tl.constexpr, - NB: tl.constexpr, - BT: tl.constexpr, -): - i_t = tl.program_id(0) - p_x = tl.make_block_ptr(x, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) - p_y = tl.make_block_ptr(y, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) - p_rstd = tl.make_block_ptr(rstd, (T,), (1,), (i_t * BT,), (BT,), (0,)) - - b_x = tl.load(p_x, boundary_check=(0, 1)).to(tl.float32) - b_rstd = 1 / tl.sqrt(tl.sum(b_x * b_x, 1) + eps) - b_y = b_x * b_rstd[:, None] - - tl.store(p_y, b_y.to(p_y.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_rstd, b_rstd.to(p_rstd.dtype.element_ty), boundary_check=(0,)) - - -@triton.autotune( - configs=[triton.Config({"BT": BT}, num_warps=num_warps) for num_warps in [1, 2, 4, 8, 16] for BT in BT_LIST], - key=["D", "NB"], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=["T"]) -def l2norm_bwd_kernel( - y, - rstd, - dy, - dx, - eps, - T, - D: tl.constexpr, - BD: tl.constexpr, - NB: tl.constexpr, - BT: tl.constexpr, -): - i_t = tl.program_id(0) - p_y = tl.make_block_ptr(y, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) - p_rstd = tl.make_block_ptr(rstd, (T,), (1,), (i_t * BT,), (BT,), (0,)) - p_dy = tl.make_block_ptr(dy, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) - p_dx = tl.make_block_ptr(dx, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) - - b_y = tl.load(p_y, boundary_check=(0, 1)).to(tl.float32) - b_rstd = tl.load(p_rstd, boundary_check=(0,)).to(tl.float32) - b_dy = tl.load(p_dy, boundary_check=(0, 1)).to(tl.float32) - b_dx = b_dy * b_rstd[:, None] - tl.sum(b_dy * b_y, 1)[:, None] * b_y * b_rstd[:, None] - tl.store(p_dx, b_dx.to(p_dx.dtype.element_ty), boundary_check=(0, 1)) - - -def l2norm_fwd( - x: torch.Tensor, - eps: float = 1e-6, - output_dtype: torch.dtype | None = None, -): - x_shape_og = x.shape - x = x.view(-1, x.shape[-1]) - # allocate output - if output_dtype is None: - y = torch.empty_like(x) - else: - y = torch.empty_like(x, dtype=output_dtype) - assert y.stride(-1) == 1 - T, D = x.shape[0], x.shape[-1] - # Less than 64KB per feature: enqueue fused kernel - MAX_FUSED_SIZE = 65536 // x.element_size() - BD = min(MAX_FUSED_SIZE, triton.next_power_of_2(D)) - if D > BD: - raise RuntimeError("This layer doesn't support feature dim >= 64KB.") - - rstd = torch.empty((T,), dtype=torch.float32, device=x.device) - if D <= 512: - # NOTE(tylerr): Avoid excessive recompilation and autotuning by tolerating a larger range - # of T before recompiling the kernel. - # NB = triton.cdiv(T, 2048) - NB = triton.cdiv(T, 2048 * 32) - - def grid(meta): - return (triton.cdiv(T, meta["BT"]),) - - l2norm_fwd_kernel[grid]( - x=x, - y=y, - rstd=rstd, - eps=eps, - T=T, - D=D, - BD=BD, - NB=NB, - ) - else: - l2norm_fwd_kernel1[(T,)]( - x=x, - y=y, - rstd=rstd, - eps=eps, - D=D, - BD=BD, - ) - return y.view(x_shape_og), rstd.view(x_shape_og[:-1]) - - -def l2norm_bwd( - y: torch.Tensor, - rstd: torch.Tensor, - dy: torch.Tensor, - eps: float = 1e-6, -): - y_shape_og = y.shape - y = y.view(-1, dy.shape[-1]) - dy = dy.view(-1, dy.shape[-1]) - assert dy.shape == y.shape - # allocate output - dx = torch.empty_like(y) - T, D = y.shape[0], y.shape[-1] - # Less than 64KB per feature: enqueue fused kernel - MAX_FUSED_SIZE = 65536 // y.element_size() - BD = min(MAX_FUSED_SIZE, triton.next_power_of_2(D)) - if D > BD: - raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") - - if D <= 512: - # NOTE(tylerr): Avoid excessive recompilation and autotuning by tolerating a larger range - # of T before recompiling the kernel. - # NB = triton.cdiv(T, 2048) - NB = triton.cdiv(T, 2048 * 32) - - def grid(meta): - return (triton.cdiv(T, meta["BT"]),) - - l2norm_bwd_kernel[grid]( - y=y, - rstd=rstd, - dy=dy, - dx=dx, - eps=eps, - T=T, - D=D, - BD=BD, - NB=NB, - ) - else: - l2norm_bwd_kernel1[(T,)]( - y=y, - rstd=rstd, - dy=dy, - dx=dx, - eps=eps, - D=D, - BD=BD, - ) - - return dx.view(y_shape_og) - - -class L2NormFunction(torch.autograd.Function): - @staticmethod - @input_guard - def forward( - ctx, - x, - eps=1e-6, - output_dtype=None, - ): - y, rstd = l2norm_fwd(x, eps, output_dtype) - ctx.eps = eps - ctx.x_dtype = x.dtype - ctx.save_for_backward(y, rstd) - return y - - @staticmethod - @input_guard - def backward(ctx, dy): - y, rstd = ctx.saved_tensors - dx = l2norm_bwd(y, rstd, dy, ctx.eps) - return dx, None, None - - -def l2norm( - x: torch.Tensor, - eps: float = 1e-6, - output_dtype: torch.dtype | None = None, -) -> torch.Tensor: - return L2NormFunction.apply(x, eps, output_dtype) diff --git a/src/xorl/ops/linear_attention/modules/layernorm.py b/src/xorl/ops/linear_attention/modules/layernorm.py deleted file mode 100644 index 5155eb8a..00000000 --- a/src/xorl/ops/linear_attention/modules/layernorm.py +++ /dev/null @@ -1,57 +0,0 @@ -from __future__ import annotations - -# 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 - - -class RMSNorm(nn.Module): - def __init__( - self, - hidden_size: int, - elementwise_affine: bool = True, - bias: bool = False, - eps: float = 1e-5, - device: torch.device | None = None, - dtype: torch.dtype | None = None, - ) -> None: - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.hidden_size = hidden_size - self.elementwise_affine = elementwise_affine - self.eps = eps - - self.register_parameter("weight", None) - self.register_parameter("bias", None) - if elementwise_affine: - self.weight = nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - if bias: - self.bias = nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - - self.reset_parameters() - - def reset_parameters(self) -> None: - if self.weight is not None: - nn.init.ones_(self.weight) - if self.bias is not None: - nn.init.zeros_(self.bias) - - def forward( - self, - x: torch.Tensor, - residual: torch.Tensor | None = None, - prenorm: bool = False, - residual_in_fp32: bool = False, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - del residual_in_fp32 - residual_out = x if residual is None else x + residual - norm_input = residual_out.float() - inv_rms = torch.rsqrt(norm_input.square().mean(dim=-1, keepdim=True) + self.eps) - y = norm_input * inv_rms - if self.weight is not None: - y = y * self.weight.float() - if self.bias is not None: - y = y + self.bias.float() - y = y.to(x.dtype) - return y if not prenorm else (y, residual_out) diff --git a/src/xorl/ops/linear_attention/ops/__init__.py b/src/xorl/ops/linear_attention/ops/__init__.py index ebe4f085..59a690a3 100644 --- a/src/xorl/ops/linear_attention/ops/__init__.py +++ b/src/xorl/ops/linear_attention/ops/__init__.py @@ -1,4 +1,4 @@ -from .gated_delta_rule import chunk_gated_delta_rule, fused_recurrent_gated_delta_rule +from fla.ops.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 deleted file mode 100644 index fcad84f5..00000000 --- a/src/xorl/ops/linear_attention/ops/common/chunk_delta_h.py +++ /dev/null @@ -1,597 +0,0 @@ -# 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 triton.language as tl - -from xorl.ops.linear_attention.ops.utils import prepare_chunk_indices, prepare_chunk_offsets -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.autotune( - configs=[ - 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]) - ], - key=["H", "K", "V", "BT", "USE_EXP2"], - use_cuda_graph=USE_CUDA_GRAPH, - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=["T"]) -def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( - k, - v, - w, - v_new, - g, - gk, - h, - h0, - ht, - cu_seqlens, - chunk_offsets, - T, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BV: tl.constexpr, - USE_G: tl.constexpr, - USE_GK: tl.constexpr, - USE_INITIAL_STATE: tl.constexpr, - STORE_FINAL_STATE: tl.constexpr, - SAVE_NEW_VALUE: tl.constexpr, - USE_EXP2: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_v, i_nh = tl.program_id(0), tl.program_id(1) - i_n, i_h = i_nh // H, i_nh % H - if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - NT = tl.cdiv(T, BT) - boh = tl.load(chunk_offsets + i_n).to(tl.int32) - else: - bos, eos = i_n * T, i_n * T + T - NT = tl.cdiv(T, BT) - boh = i_n * NT - - # [BK, BV] - b_h1 = tl.zeros([64, BV], dtype=tl.float32) - if K > 64: - b_h2 = tl.zeros([64, BV], dtype=tl.float32) - if K > 128: - b_h3 = tl.zeros([64, BV], dtype=tl.float32) - if K > 192: - b_h4 = tl.zeros([64, BV], dtype=tl.float32) - - # calculate offset - 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 - if SAVE_NEW_VALUE: - v_new += (bos * H + i_h).to(tl.int64) * V - - if USE_INITIAL_STATE: - h0 = h0 + i_nh * K * V - if STORE_FINAL_STATE: - ht = ht + i_nh * K * V - - # load initial state - if USE_INITIAL_STATE: - p_h0_1 = tl.make_block_ptr(h0, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) - b_h1 += tl.load(p_h0_1, boundary_check=(0, 1)).to(tl.float32) - if K > 64: - p_h0_2 = tl.make_block_ptr(h0, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0)) - b_h2 += tl.load(p_h0_2, boundary_check=(0, 1)).to(tl.float32) - if K > 128: - p_h0_3 = tl.make_block_ptr(h0, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0)) - b_h3 += tl.load(p_h0_3, boundary_check=(0, 1)).to(tl.float32) - if K > 192: - p_h0_4 = tl.make_block_ptr(h0, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0)) - b_h4 += tl.load(p_h0_4, boundary_check=(0, 1)).to(tl.float32) - - # 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)) - 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)) - 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)) - 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)) - 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)) - 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)) - 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)) - 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)) - 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)) - 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)) - 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 - if USE_G: - m_t = (i_t * BT + tl.arange(0, BT)) < T - b_g_last = tl.load(g + (bos * H + last_idx * H + i_h).to(tl.int64)).to(tl.float32) - p_g = tl.make_block_ptr(g + (bos * H + i_h).to(tl.int64), (T,), (H,), (i_t * BT,), (BT,), (0,)) - b_g = tl.load(p_g, boundary_check=(0,)).to(tl.float32) - if USE_EXP2: - b_v = b_v * tl.where(m_t, exp2(b_g_last - b_g), 0)[:, None] - b_g_last = exp2(b_g_last) - else: - b_v = b_v * tl.where(m_t, exp(b_g_last - b_g), 0)[:, None] - b_g_last = exp(b_g_last) - b_h1 *= b_g_last - if K > 64: - b_h2 *= b_g_last - if K > 128: - b_h3 *= b_g_last - if K > 192: - b_h4 *= b_g_last - - 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.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.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.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.0).to( - tl.float32 - ) - if USE_EXP2: - b_h4 *= exp2(b_gk_last4)[:, None] - else: - b_h4 *= exp(b_gk_last4)[:, None] - - 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)) - 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)) - 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)) - 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)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_h4 += tl.dot(b_k, b_v) - - if STORE_FINAL_STATE: - p_ht = tl.make_block_ptr(ht, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) - tl.store(p_ht, b_h1.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) - if K > 64: - p_ht = tl.make_block_ptr(ht, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0)) - tl.store(p_ht, b_h2.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) - if K > 128: - p_ht = tl.make_block_ptr(ht, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0)) - tl.store(p_ht, b_h3.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) - if K > 192: - p_ht = tl.make_block_ptr(ht, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0)) - 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.autotune( - configs=[ - 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]) - ], - 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"]) -def chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64( - q, - k, - w, - g, - gk, - dht, - dh0, - do, - dh, - dv, - dv2, - cu_seqlens, - chunk_offsets, - scale, - T, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BV: tl.constexpr, - USE_G: tl.constexpr, - USE_GK: tl.constexpr, - USE_INITIAL_STATE: tl.constexpr, - USE_FINAL_STATE_GRADIENT: tl.constexpr, - USE_EXP2: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_v, i_nh = tl.program_id(0), tl.program_id(1) - i_n, i_h = i_nh // H, i_nh % H - if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - NT = tl.cdiv(T, BT) - boh = tl.load(chunk_offsets + i_n).to(tl.int32) - else: - bos, eos = i_n * T, i_n * T + T - NT = tl.cdiv(T, BT) - boh = i_n * NT - - # [BK, BV] - b_dh1 = tl.zeros([64, BV], dtype=tl.float32) - if K > 64: - b_dh2 = tl.zeros([64, BV], dtype=tl.float32) - if K > 128: - b_dh3 = tl.zeros([64, BV], dtype=tl.float32) - if K > 192: - b_dh4 = tl.zeros([64, BV], dtype=tl.float32) - - # calculate offset - q += (bos * H + i_h).to(tl.int64) * K - k += (bos * H + i_h).to(tl.int64) * K - w += (bos * H + i_h).to(tl.int64) * K - 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 - if USE_GK: - gk += (bos * H + i_h).to(tl.int64) * K - - if USE_INITIAL_STATE: - dh0 += i_nh * K * V - if USE_FINAL_STATE_GRADIENT: - 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)) - b_dh1 += tl.load(p_dht1, boundary_check=(0, 1)) - if K > 64: - p_dht2 = tl.make_block_ptr(dht, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0)) - b_dh2 += tl.load(p_dht2, boundary_check=(0, 1)) - if K > 128: - p_dht3 = tl.make_block_ptr(dht, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0)) - b_dh3 += tl.load(p_dht3, boundary_check=(0, 1)) - if K > 192: - p_dht4 = tl.make_block_ptr(dht, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0)) - b_dh4 += tl.load(p_dht4, boundary_check=(0, 1)) - - 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)) - 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)) - 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)) - 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)) - 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 - if USE_G: - bg_last = tl.load(g + (bos + last_idx) * H + i_h).to(tl.float32) - 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,)).to(tl.float32) - if USE_EXP2: - bg_last_exp = exp2(bg_last) - b_g_exp = exp2(b_g) - else: - 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)) - - 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)) - 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.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)) - 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.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)) - 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.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)) - 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.0).to(tl.float32) - b_dv += tl.dot(b_k, b_dh4.to(b_k.dtype)) - - if USE_G: - m_t = (i_t * BT + tl.arange(0, BT)) < T - if USE_EXP2: - b_dv *= tl.where(m_t, exp2(bg_last - b_g), 0)[:, None] - else: - b_dv *= tl.where(m_t, exp(bg_last - b_g), 0)[:, None] - b_dv += tl.load(p_dv, boundary_check=(0, 1)) - - 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)) - b_w = tl.load(p_w, boundary_check=(0, 1)) - b_q = tl.load(p_q, boundary_check=(0, 1)) - if USE_G: - b_dh1 *= bg_last_exp - b_q = b_q * b_g_exp[None, :] - if USE_GK: - if USE_EXP2: - b_dh1 *= exp2(b_gk_last1[:, None]) - else: - 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)) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_w = tl.load(p_w, boundary_check=(0, 1)) - if USE_G: - b_dh2 *= bg_last_exp - b_q = b_q * b_g_exp[None, :] - if USE_GK: - if USE_EXP2: - b_dh2 *= exp2(b_gk_last2[:, None]) - else: - 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)) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_w = tl.load(p_w, boundary_check=(0, 1)) - if USE_G: - b_dh3 *= bg_last_exp - b_q = b_q * b_g_exp[None, :] - if USE_GK: - if USE_EXP2: - b_dh3 *= exp2(b_gk_last3[:, None]) - else: - 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)) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_w = tl.load(p_w, boundary_check=(0, 1)) - if USE_G: - b_dh4 *= bg_last_exp - b_q = b_q * b_g_exp[None, :] - if USE_GK: - if USE_EXP2: - b_dh4 *= exp2(b_gk_last4[:, None]) - else: - b_dh4 *= exp(b_gk_last4[:, None]) - b_dh4 += 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 USE_INITIAL_STATE: - p_dh0 = tl.make_block_ptr(dh0, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) - tl.store(p_dh0, b_dh1.to(p_dh0.dtype.element_ty), boundary_check=(0, 1)) - if K > 64: - p_dh1 = tl.make_block_ptr(dh0, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0)) - tl.store(p_dh1, b_dh2.to(p_dh1.dtype.element_ty), boundary_check=(0, 1)) - if K > 128: - p_dh2 = tl.make_block_ptr(dh0, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0)) - tl.store(p_dh2, b_dh3.to(p_dh2.dtype.element_ty), boundary_check=(0, 1)) - if K > 192: - p_dh3 = tl.make_block_ptr(dh0, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0)) - tl.store(p_dh3, b_dh4.to(p_dh3.dtype.element_ty), boundary_check=(0, 1)) - - -def chunk_gated_delta_rule_fwd_h( - k: torch.Tensor, - w: torch.Tensor, - u: torch.Tensor, - g: torch.Tensor | None = None, - gk: torch.Tensor | None = None, - initial_state: torch.Tensor | None = None, - output_final_state: bool = False, - chunk_size: int = 64, # SY: remove this argument and force chunk size 64? - save_new_value: bool = True, - cu_seqlens: torch.LongTensor | None = None, - cu_seqlens_cpu: torch.LongTensor | None = None, - chunk_indices: torch.LongTensor | None = None, - use_exp2: bool = False, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: - B, T, H, K, V = *k.shape, u.shape[-1] - BT = chunk_size - - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) - # N: the actual number of sequences in the batch with either equal or variable lengths - if cu_seqlens is None: - N, NT, chunk_offsets = B, triton.cdiv(T, BT), None - else: - N, NT, chunk_offsets = len(cu_seqlens) - 1, len(chunk_indices), prepare_chunk_offsets(cu_seqlens, BT) - assert K <= 256, "current kernel does not support head dimension larger than 256." - - h = k.new_empty(B, NT, H, K, V) - # Ensure final output is zeros - # vLLM will use padding for CUDA Graph - 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) - - chunk_gated_delta_rule_fwd_kernel_h_blockdim64[grid]( - k=k, - v=u, - w=w, - v_new=v_new, - g=g, - gk=gk, - h=h, - h0=initial_state, - ht=final_state, - cu_seqlens=cu_seqlens, - chunk_offsets=chunk_offsets, - T=T, - H=H, - K=K, - V=V, - BT=BT, - USE_EXP2=use_exp2, - ) - return h, v_new, final_state - - -def chunk_gated_delta_rule_bwd_dhu( - q: torch.Tensor, - k: torch.Tensor, - w: torch.Tensor, - do: torch.Tensor, - dv: torch.Tensor, - g: torch.Tensor | None = None, - gk: torch.Tensor | None = None, - h0: torch.Tensor | None = None, - dht: torch.Tensor | None = None, - scale: float | None = None, - cu_seqlens: torch.LongTensor | None = None, - chunk_size: int = 64, # SY: remove this argument and force chunk size 64? - chunk_indices: torch.LongTensor | None = None, - use_exp2: bool = False, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - B, T, H, K, V = *q.shape, do.shape[-1] - # N: the actual number of sequences in the batch with either equal or variable lengths - BT = 64 - assert K <= 256, "current kernel does not support head dimension being larger than 256." - - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) - if cu_seqlens is None: - N, NT, chunk_offsets = B, triton.cdiv(T, BT), None - else: - N, NT, chunk_offsets = len(cu_seqlens) - 1, len(chunk_indices), prepare_chunk_offsets(cu_seqlens, BT) - - dh = q.new_empty(B, NT, H, K, V) - 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) - - chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64[grid]( - q=q, - k=k, - w=w, - g=g, - gk=gk, - dht=dht, - dh0=dh0, - do=do, - dh=dh, - dv=dv, - dv2=dv2, - cu_seqlens=cu_seqlens, - chunk_offsets=chunk_offsets, - scale=scale, - T=T, - H=H, - K=K, - V=V, - BT=BT, - USE_EXP2=use_exp2, - ) - return dh, dh0, dv2 diff --git a/src/xorl/ops/linear_attention/ops/common/chunk_o.py b/src/xorl/ops/linear_attention/ops/common/chunk_o.py deleted file mode 100644 index 4a58b2a6..00000000 --- a/src/xorl/ops/linear_attention/ops/common/chunk_o.py +++ /dev/null @@ -1,707 +0,0 @@ -# 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 triton.language as tl - -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_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]) -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.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), - ], - key=["H", "K", "V", "BT"], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=["T"]) -def chunk_fwd_kernel_o( - q, - k, - v, - h, - g, - g_gamma, - o, - cu_seqlens, - chunk_indices, - scale, - T, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - USE_G: tl.constexpr, - USE_G_GAMMA: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_b, i_h = i_bh // H, i_bh % H - - if IS_VARLEN: - i_tg = i_t - i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - NT = tl.cdiv(T, BT) - else: - NT = tl.cdiv(T, BT) - i_tg = i_b * NT + i_t - bos, eos = i_b * T, i_b * T + T - - # offset calculation - q += (bos * H + i_h) * K - 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 - - 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_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)) - # [BK, BT] - b_k = tl.load(p_k, boundary_check=(0, 1)) - # [BK, BV] - b_h = tl.load(p_h, boundary_check=(0, 1)) - - # [BT, BK] @ [BK, BV] -> [BT, BV] - b_o += tl.dot(b_q, b_h) - # [BT, BK] @ [BK, BT] -> [BT, BT] - b_A += tl.dot(b_q, b_k) - - if USE_G: - g += bos * H + i_h - p_g = tl.make_block_ptr(g, (T,), (H,), (i_t * BT,), (BT,), (0,)) - b_g = tl.load(p_g, boundary_check=(0,)) - b_o = b_o * exp(b_g)[:, None] - b_A = b_A * exp(b_g[:, None] - b_g[None, :]) - - if USE_G_GAMMA: - b_gamma = tl.load(g_gamma + i_h) - b_g = b_gamma * (tl.arange(0, BT) + 1) - b_o = b_o * exp(b_g)[:, None] - b_A = b_A * exp(b_g[:, None] - b_g[None, :]) - - o_t = i_t * BT + tl.arange(0, BT) - m_t = o_t < T - 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)) - - b_v = tl.load(p_v, boundary_check=(0, 1)) - # to fix mma -> mma layout conversion - # already solved by triton v3.2 or higher - b_o = b_o * scale + tl.dot(b_A.to(b_v.dtype), b_v) * scale - 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.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"], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=["T"]) -def chunk_bwd_kernel_dqkwg( - q, - k, - v, - g, - g_gamma, - h, - do, - dh, - dq, - dk, - dw, - dv, - dg, - cu_seqlens, - chunk_indices, - scale, - B: tl.constexpr, - T, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - USE_G: tl.constexpr, - USE_G_GAMMA: tl.constexpr, - USE_DW: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_k, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_b, i_h = i_bh // H, i_bh % H - - all = B * T - if IS_VARLEN: - i_tg = i_t - i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - NT = tl.cdiv(T, BT) - else: - NT = tl.cdiv(T, BT) - i_tg = i_b * NT + i_t - bos, eos = i_b * T, i_b * T + T - - # 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 - q += (bos * H + i_h) * K - k += (bos * H + i_h) * K - dq += (bos * H + i_h) * K - dk += (bos * H + i_h) * K - - # for delta rule only - if USE_DW: - dw += (bos * H + i_h) * K - dv += (bos * H + i_h) * V - - if USE_G: - dg += i_k * all * H - b_dg_last = tl.zeros([1], dtype=tl.float32) if USE_G else None - if USE_G_GAMMA: - b_gamma = tl.load(g_gamma + i_h) - b_g = b_gamma * (tl.arange(0, BT) + 1) - b_g_last = b_gamma * min(BT, T - i_t * BT) - b_dq = tl.zeros([BT, BK], dtype=tl.float32) - b_dk = tl.zeros([BT, BK], dtype=tl.float32) - b_ds = tl.zeros([BT, BT], dtype=tl.float32) - 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_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] - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_do = tl.load(p_do, boundary_check=(0, 1)) - # [BV, BK] - 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) - # [BT, BV] @ [BV, BT] -> [BT, BT] - b_ds += tl.dot(b_do, tl.trans(b_v)) - # [BT, BV] @ [BV, BK] -> [BT, BK] - b_dq += tl.dot(b_do, b_h.to(b_do.dtype)) - # [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)) - 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)) - 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)) - 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)) - - o_t = i_t * BT + tl.arange(0, BT) - m_t = o_t < T - m_A = (o_t[:, None] >= o_t[None, :]) & (m_t[:, None] & m_t) - if USE_G: - b_dg = tl.zeros([BT], dtype=tl.float32) - g += bos * H + i_h - dg += bos * H + i_h - p_g = tl.make_block_ptr(g, (T,), (H,), (i_t * BT,), (BT,), (0,)) - b_g = tl.load(p_g, boundary_check=(0,)) - b_g_last = tl.load(g + (min(i_t * BT + BT, T) - 1) * H) - b_dg_last *= exp(b_g_last) - - b_dq = b_dq * exp(b_g)[:, None] * scale - b_dg += tl.sum(b_dq * b_q, axis=1) - - b_dk = b_dk * tl.where(m_t, exp(-b_g + b_g_last), 0)[:, None] - b_dg -= tl.sum(b_k * b_dk, axis=1) - b_dg_last += tl.sum(b_dk * b_k) - - b_ds = tl.where(m_A, b_ds * exp(b_g[:, None] - b_g[None, :]), 0) * scale - b_ds2 = b_ds * tl.dot(b_q, tl.trans(b_k)) - b_dg += tl.sum(b_ds2, axis=1) - b_dg -= tl.sum(b_ds2, axis=0) - - b_ds = b_ds.to(b_k.dtype) - # [BT, BK] - b_dq += tl.dot(b_ds, b_k) - b_dk += tl.dot(tl.trans(b_ds), b_q) - p_dg = tl.make_block_ptr(dg, (T,), (H,), (i_t * BT,), (BT,), (0,)) - # (SY 09/21) revcumsum in a separate kernel due to strange triton compiler issue - # b_dg = tl.dot(tl.where(o_t[:, None] <= o_t[None, :], 1., 0.), b_dg, allow_tf32=False) + b_dg_last) - b_dg = tl.where(o_t < min(i_t * BT + BT, T) - 1, b_dg, b_dg + b_dg_last) - tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0,)) - - elif USE_G_GAMMA: - b_dq = b_dq * exp(b_g)[:, None] * scale - b_dk = b_dk * tl.where(m_t, exp(-b_g + b_g_last), 0)[:, None] - b_ds = tl.where(m_A, b_ds * exp(b_g[:, None] - b_g[None, :]), 0) * scale - b_ds = b_ds.to(b_k.dtype) - # [BT, BK] - b_dq += tl.dot(b_ds, b_k) - b_dk += tl.dot(tl.trans(b_ds), b_q) - tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) - - else: - b_ds = tl.where(m_A, b_ds, 0) - b_ds = b_ds.to(b_k.dtype) - b_dq += tl.dot(b_ds, b_k) - b_dk += tl.dot(tl.trans(b_ds), b_q) * scale - b_dq *= scale - tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) - 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.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"], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=["T"]) -def chunk_bwd_kernel_dv( - q, - k, - g, - g_gamma, - do, - dv, - dh, - cu_seqlens, - chunk_indices, - scale, - T, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - USE_G: tl.constexpr, - USE_G_GAMMA: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_b, i_h = i_bh // H, i_bh % H - if IS_VARLEN: - i_tg = i_t - i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - NT = tl.cdiv(T, BT) - else: - NT = tl.cdiv(T, BT) - i_tg = i_b * NT + i_t - bos, eos = i_b * T, i_b * T + T - - b_dv = tl.zeros([BT, BV], dtype=tl.float32) - - # offset calculation - q += (bos * H + i_h) * K - 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 - - 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)) - 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) - p_dh = tl.make_block_ptr(dh, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - b_dh = tl.load(p_dh, boundary_check=(0, 1)) - b_dv += tl.dot(b_k, b_dh.to(b_k.dtype)) - - o_t = i_t * BT + tl.arange(0, BT) - m_t = o_t < T - if USE_G: - g += bos * H + i_h - p_g = tl.make_block_ptr(g, (T,), (H,), (i_t * BT,), (BT,), (0,)) - b_g = tl.load(p_g, boundary_check=(0,)) - b_g_last = tl.load(g + (min(i_t * BT + BT, T) - 1) * H) - if USE_G_GAMMA: - b_gamma = tl.load(g_gamma + i_h) - b_g = b_gamma * (tl.arange(0, BT) + 1) - b_g_last = b_gamma * min(BT, T - i_t * BT) - - m_A = (o_t[:, None] <= o_t[None, :]) & (m_t[:, None] & m_t) - if USE_G or USE_G_GAMMA: - b_A = tl.where(m_A, b_A * exp(b_g[None, :] - b_g[:, None]) * scale, 0).to(do.dtype.element_ty) - 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)) - 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.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"], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=["T"]) -def chunk_bwd_kernel_dv_local( - q, - k, - g, - g_gamma, - A, - do, - dv, - cu_seqlens, - chunk_indices, - scale, - T, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - USE_G: tl.constexpr, - USE_G_GAMMA: tl.constexpr, - USE_A: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_t, i_bh = tl.program_id(0), tl.program_id(1) - i_b, i_h = i_bh // H, i_bh % H - if IS_VARLEN: - i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - else: - bos, eos = i_b * T, i_b * T + T - - # offset calculation - q += (bos * H + i_h) * K - k += (bos * H + i_h) * K - do += (bos * H + i_h) * V - 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)) - b_A = tl.load(p_A, boundary_check=(0, 1)) - else: - if USE_G: - g += bos * H + i_h - p_g = tl.make_block_ptr(g, (T,), (H,), (i_t * BT,), (BT,), (0,)) - b_g = tl.load(p_g, boundary_check=(0,)) - if USE_G_GAMMA: - b_gamma = tl.load(g_gamma + i_h) - b_g = b_gamma * (tl.arange(0, BT) + 1) - - 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)) - - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_A += tl.dot(b_k, b_q) * scale - if USE_G or USE_G_GAMMA: - b_A *= exp(b_g[None, :] - b_g[:, None]) - - o_t = i_t * BT + tl.arange(0, BT) - m_t = o_t < T - m_A = (o_t[:, None] <= o_t[None, :]) & (m_t[:, None] & m_t) - 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)) - 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)) - - -def chunk_fwd_o( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - h: torch.Tensor, - g: torch.Tensor | None = None, - g_gamma: torch.Tensor | None = None, - scale: float | None = None, - cu_seqlens: torch.LongTensor | None = None, - chunk_size: int = 64, -) -> torch.Tensor: - B, T, H, K, V = *q.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 scale is None: - scale = k.shape[-1] ** -0.5 - - o = torch.empty_like(v) - - def grid(meta): - return (triton.cdiv(V, meta["BV"]), NT, B * H) - - chunk_fwd_kernel_o[grid]( - q=q, - k=k, - v=v, - h=h, - g=g, - g_gamma=g_gamma, - o=o, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - scale=scale, - T=T, - H=H, - K=K, - V=V, - BT=BT, - ) - return o - - -def chunk_bwd_dv( - q: torch.Tensor, - k: torch.Tensor, - do: torch.Tensor, - dh: torch.Tensor, - g: torch.Tensor | None = None, - g_gamma: torch.Tensor | None = None, - scale: float | None = None, - cu_seqlens: torch.LongTensor | None = None, - chunk_size: int = 64, -) -> torch.Tensor: - B, T, H, K, V = *k.shape, do.shape[-1] - 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): - CONST_TILING = 128 - elif check_shared_mem("ada", k.device.index): - CONST_TILING = 64 - else: - CONST_TILING = 32 - BK = min(max(triton.next_power_of_2(K), 16), CONST_TILING) - BV = min(max(triton.next_power_of_2(V), 16), CONST_TILING) - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - NV = triton.cdiv(V, BV) - if scale is None: - scale = k.shape[-1] ** -0.5 - - dv = torch.empty_like(do) - grid = (NV, NT, B * H) - chunk_bwd_kernel_dv[grid]( - q=q, - k=k, - g=g, - g_gamma=g_gamma, - do=do, - dv=dv, - dh=dh, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - scale=scale, - T=T, - H=H, - K=K, - V=V, - BT=BT, - BK=BK, - BV=BV, - ) - return dv - - -def chunk_bwd_dv_local( - q: torch.Tensor, - k: torch.Tensor, - do: torch.Tensor, - g: torch.Tensor | None = None, - g_gamma: torch.Tensor | None = None, - A: torch.Tensor | None = None, - scale: float = None, - cu_seqlens: torch.LongTensor | None = None, - chunk_size: int = 64, - chunk_indices: torch.LongTensor | None = None, -) -> torch.Tensor: - B, T, H, K, V = *k.shape, do.shape[-1] - BT = chunk_size - 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): - CONST_TILING = 128 - elif check_shared_mem("ada", k.device.index): - CONST_TILING = 64 - else: - CONST_TILING = 32 - BK = min(max(triton.next_power_of_2(K), 16), CONST_TILING) - BV = min(max(triton.next_power_of_2(V), 16), CONST_TILING) - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - - dv = torch.empty_like(do) - grid = (NT, B * H) - chunk_bwd_kernel_dv_local[grid]( - q=q, - k=k, - g=g, - g_gamma=g_gamma, - A=A, - do=do, - dv=dv, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - scale=scale, - T=T, - H=H, - K=K, - V=V, - BT=BT, - BK=BK, - BV=BV, - ) - return dv - - -def chunk_bwd_dqkwg( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - do: torch.Tensor, - h: torch.Tensor, - dh: torch.Tensor, - w: torch.Tensor | None = None, - g: torch.Tensor | None = None, - g_gamma: torch.Tensor | None = None, - dv: torch.Tensor | None = None, - scale: float | None = None, - 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): - CONST_TILING = 128 - elif check_shared_mem("ada", k.device.index): - CONST_TILING = 64 - else: - CONST_TILING = 32 - BK = min(max(triton.next_power_of_2(K), 16), CONST_TILING) - BV = min(max(triton.next_power_of_2(V), 16), CONST_TILING) - NK = triton.cdiv(K, BK) - dq = torch.empty_like(q) - dk = torch.empty_like(k) - dg = torch.empty(NK, *g.shape, dtype=torch.float32, device=g.device) if g is not None else None - dw = torch.empty_like(w) if w is not None else None - - grid = (NK, NT, B * H) - chunk_bwd_kernel_dqkwg[grid]( - q=q, - k=k, - v=v, - g=g, - g_gamma=g_gamma, - h=h, - do=do, - dh=dh, - dw=dw, - dq=dq, - dk=dk, - dv=dv, - dg=dg, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - scale=scale, - B=B, - T=T, - H=H, - K=K, - V=V, - BT=BT, - BK=BK, - BV=BV, - ) - - if dg is not None: - dg = dg.sum(0) - return dq, dk, dw, dg 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 deleted file mode 100644 index 62cf89e9..00000000 --- a/src/xorl/ops/linear_attention/ops/common/chunk_scaled_dot_kkt.py +++ /dev/null @@ -1,127 +0,0 @@ -# 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 triton.language as tl - -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 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.autotune( - configs=[ - 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"], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=["T"]) -def chunk_scaled_dot_kkt_fwd_kernel( - k, - g, - beta, - A, - cu_seqlens, - chunk_indices, - T, - H: tl.constexpr, - K: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - IS_VARLEN: tl.constexpr, - USE_G: tl.constexpr, -): - i_t, i_bh = tl.program_id(0), tl.program_id(1) - i_b, i_h = i_bh // H, i_bh % H - if IS_VARLEN: - i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - else: - bos, eos = i_b * T, i_b * T + T - 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,)) - 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)) - 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,)) - b_g = tl.load(p_g, boundary_check=(0,)) - b_g_diff = b_g[:, None] - b_g[None, :] - b_A *= exp(b_g_diff) - b_A *= b_b[:, None] - - 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)) - tl.store(p_A, b_A.to(p_A.dtype.element_ty), boundary_check=(0, 1)) - - -def chunk_scaled_dot_kkt_fwd( - k: torch.Tensor, - g: torch.Tensor | None = None, - beta: torch.Tensor | None = None, - cu_seqlens: torch.LongTensor | None = None, - chunk_size: int = 64, - output_dtype: torch.dtype = torch.float32, -) -> torch.Tensor: - r""" - Compute beta * K * K^T. - - Args: - k (torch.Tensor): - The key tensor of shape `[B, T, H, K]`. - beta (torch.Tensor): - The beta tensor of shape `[B, T, H]`. - g (torch.Tensor): - The cumulative sum of the gate tensor of shape `[B, T, H]`. Default: `None`. - gk (torch.Tensor): - The cumulative sum of the gate tensor of shape `[B, T, H, K]` applied to the key tensor. Default: `None`. - cu_seqlens (torch.LongTensor): - The cumulative sequence lengths of the input tensor. - Default: None - chunk_size (int): - The chunk size. Default: 64. - output_dtype (torch.dtype): - The dtype of the output tensor. Default: `torch.float32` - - Returns: - beta * K * K^T of shape `[B, T, H, BT]` where `BT` is the chunk size. - """ - B, T, H, K = k.shape - 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) - A = torch.empty(B, T, H, BT, device=k.device, dtype=output_dtype) - chunk_scaled_dot_kkt_fwd_kernel[(NT, B * H)]( - k=k, - g=g, - beta=beta, - A=A, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - T=T, - H=H, - K=K, - BT=BT, - ) - return A 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 deleted file mode 100644 index 59bb550c..00000000 --- a/src/xorl/ops/linear_attention/ops/cp/chunk_delta_h.py +++ /dev/null @@ -1,1233 +0,0 @@ -# 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. - -from __future__ import annotations - -from typing import TYPE_CHECKING - -import torch -import torch.distributed as dist -import triton -import triton.language as tl - -from xorl.ops.linear_attention.ops.cp.comm import all_gather_into_tensor -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.autotune( - configs=[ - 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"], - use_cuda_graph=USE_CUDA_GRAPH, - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=["T"]) -def pre_process_fwd_kernel_stage1( - k, - v, - w, - g, - gk, - hm, - cu_seqlens, - T, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BV: tl.constexpr, - USE_G: tl.constexpr, - USE_GK: tl.constexpr, - USE_EXP2: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_v, i_h = tl.program_id(0), tl.program_id(1) - i_n = 0 - if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) - T = (eos - bos).to(tl.int32) - NT = tl.cdiv(T, BT) - else: - bos, eos = (i_n * T).to(tl.int64), (i_n * T + T).to(tl.int64) - NT = tl.cdiv(T, BT) - - hm += i_h * K * (K + V) - v += ((bos * H + i_h) * V).to(tl.int64) - k += ((bos * H + i_h) * K).to(tl.int64) - w += ((bos * H + i_h) * K).to(tl.int64) - stride_v = H * V - stride_k = H * K - - b_h1 = tl.zeros([64, BV], dtype=tl.float32) - if K > 64: - b_h2 = tl.zeros([64, BV], dtype=tl.float32) - if K > 128: - b_h3 = tl.zeros([64, BV], dtype=tl.float32) - if K > 192: - b_h4 = tl.zeros([64, BV], dtype=tl.float32) - - for i_t in range(NT): - p_w = tl.make_block_ptr(w, (T, K), (stride_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), (stride_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), (stride_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), (stride_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), (stride_v, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_v - - last_idx = min((i_t + 1) * BT, T) - 1 - if USE_G: - m_t = (i_t * BT + tl.arange(0, BT)) < T - b_g_last = tl.load(g + bos * H + last_idx * H + i_h).to(tl.float32) - 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,)).to(tl.float32) - if USE_EXP2: - b_v = b_v * tl.where(m_t, exp2(b_g_last - b_g), 0)[:, None] - b_g_last = exp2(b_g_last) - else: - b_v = b_v * tl.where(m_t, exp(b_g_last - b_g), 0)[:, None] - b_g_last = exp(b_g_last) - b_h1 *= b_g_last - if K > 64: - b_h2 *= b_g_last - if K > 128: - b_h3 *= b_g_last - if K > 192: - b_h4 *= b_g_last - - 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.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.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.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.0).to( - tl.float32 - ) - if USE_EXP2: - b_h4 *= exp2(b_gk_last4)[:, None] - else: - b_h4 *= exp(b_gk_last4)[:, None] - - b_v = b_v.to(k.dtype.element_ty) - - p_k = tl.make_block_ptr(k, (K, T), (1, stride_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, stride_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, stride_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, stride_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) - - p_h1 = tl.make_block_ptr(hm, (K, V), (K + 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(hm, (K, V), (K + 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(hm, (K, V), (K + 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(hm, (K, V), (K + 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)) - - -@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) - for num_warps in [2, 4] - for num_stages in [2, 3, 4] - for BK2 in [32] - ], - key=["H", "BT", "USE_EXP2", "FORWARD"], - use_cuda_graph=USE_CUDA_GRAPH, - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=["T"]) -def pre_process_fwd_bwd_kernel_stage2( - k, - w, - g, - gk, - hm, - cu_seqlens, - T, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - USE_G: tl.constexpr, - USE_GK: tl.constexpr, - USE_EXP2: tl.constexpr, - IS_VARLEN: tl.constexpr, - BK1: tl.constexpr, - BK2: tl.constexpr, - FORWARD: tl.constexpr = True, -): - i_k_col, i_h = tl.program_id(0), tl.program_id(1) - i_n = 0 - if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - NT = tl.cdiv(T, BT) - else: - bos, eos = i_n * T, i_n * T + T - NT = tl.cdiv(T, BT) - - hm += i_h * K * (K + V) - k += ((bos * H + i_h) * K).to(tl.int64) - w += ((bos * H + i_h) * K).to(tl.int64) - stride_k = H * K - - row = tl.arange(0, BK1) - col = tl.arange(0, BK2) + i_k_col * BK2 - - b_m = tl.where(row[:, None] == col[None, :], 1.0, 0.0) - for _i_t in range(NT): - if FORWARD: - i_t = _i_t - else: - i_t = NT - 1 - _i_t - p_k = tl.make_block_ptr(k, (T, K), (stride_k, 1), (i_t * BT, 0), (BT, BK1), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - p_w = tl.make_block_ptr(w, (T, K), (stride_k, 1), (i_t * BT, 0), (BT, BK1), (1, 0)) - b_w = tl.load(p_w, boundary_check=(0, 1)) - last_idx = min((i_t + 1) * BT, T) - 1 - if USE_G: - m_t = (i_t * BT + tl.arange(0, BT)) < T - b_g_last = tl.load(g + bos * H + last_idx * H + i_h).to(tl.float32) - 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,)).to(tl.float32) - if USE_EXP2: - b_k = b_k * tl.where(m_t, exp2(b_g_last - b_g), 0)[:, None] - b_g_last = exp2(b_g_last) - else: - b_k = b_k * tl.where(m_t, exp(b_g_last - b_g), 0)[:, None] - 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.0).to(tl.float32) - if USE_EXP2: - b_gk_last = exp2(b_gk_last) - else: - b_gk_last = exp(b_gk_last) - b_diag = tl.where(row[:, None] == row[None, :], b_gk_last[:, None], 0.0) - else: - b_diag = tl.where(row[:, None] == row[None, :], 1.0, 0.0) - if FORWARD: - b_kw = tl.dot(tl.trans(b_k.to(b_w.dtype)), b_w) - else: - b_kw = tl.dot(tl.trans(b_w), b_k.to(b_w.dtype)) - b_m_i = b_diag - b_kw - b_m = tl.dot(b_m_i.to(b_w.dtype), b_m.to(b_w.dtype)) - p_m = tl.make_block_ptr(hm + V, (K, K), (K + V, 1), (0, i_k_col * BK2), (BK1, BK2), (1, 0)) - 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.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"], - use_cuda_graph=USE_CUDA_GRAPH, - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=["T"]) -def pre_process_fwd_kernel_merged( - k, - v, - w, - g, - gk, - hm, - cu_seqlens, - T, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BLOCK_SIZE: tl.constexpr, - BK1: tl.constexpr, - USE_G: tl.constexpr, - USE_GK: tl.constexpr, - USE_EXP2: tl.constexpr, - IS_VARLEN: tl.constexpr, - MULTI_SEQS: tl.constexpr, -): - i_col, i_h = tl.program_id(0), tl.program_id(1) - if MULTI_SEQS: - i_n = tl.program_id(2) - hm += i_n * H * K * (K + V) + i_h * K * (K + V) - else: - i_n = 0 - hm += i_h * K * (K + V) - if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) - T = (eos - bos).to(tl.int32) - NT = tl.cdiv(T, BT) - else: - bos, eos = (i_n * T).to(tl.int64), (i_n * T + T).to(tl.int64) - NT = tl.cdiv(T, BT) - - is_h_part = i_col * BLOCK_SIZE < V - k += ((bos * H + i_h) * K).to(tl.int64) - w += ((bos * H + i_h) * K).to(tl.int64) - stride_k = H * K - - if is_h_part: - v += ((bos * H + i_h) * V).to(tl.int64) - stride_v = H * V - i_v = i_col - - b_h1 = tl.zeros([64, BLOCK_SIZE], dtype=tl.float32) - if K > 64: - b_h2 = tl.zeros([64, BLOCK_SIZE], dtype=tl.float32) - if K > 128: - b_h3 = tl.zeros([64, BLOCK_SIZE], dtype=tl.float32) - if K > 192: - b_h4 = tl.zeros([64, BLOCK_SIZE], dtype=tl.float32) - - for i_t in range(NT): - p_w = tl.make_block_ptr(w, (T, K), (stride_k, 1), (i_t * BT, 0), (BT, 64), (1, 0)) - b_w = tl.load(p_w, boundary_check=(0, 1)) - b_v_decay = tl.dot(b_w, b_h1.to(b_w.dtype)) - if K > 64: - p_w = tl.make_block_ptr(w, (T, K), (stride_k, 1), (i_t * BT, 64), (BT, 64), (1, 0)) - b_w = tl.load(p_w, boundary_check=(0, 1)) - b_v_decay += tl.dot(b_w, b_h2.to(b_w.dtype)) - if K > 128: - p_w = tl.make_block_ptr(w, (T, K), (stride_k, 1), (i_t * BT, 128), (BT, 64), (1, 0)) - b_w = tl.load(p_w, boundary_check=(0, 1)) - b_v_decay += tl.dot(b_w, b_h3.to(b_w.dtype)) - if K > 192: - p_w = tl.make_block_ptr(w, (T, K), (stride_k, 1), (i_t * BT, 192), (BT, 64), (1, 0)) - b_w = tl.load(p_w, boundary_check=(0, 1)) - b_v_decay += tl.dot(b_w, b_h4.to(b_w.dtype)) - - p_v = tl.make_block_ptr(v, (T, V), (stride_v, 1), (i_t * BT, i_v * BLOCK_SIZE), (BT, BLOCK_SIZE), (1, 0)) - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_v_decay - - last_idx = min((i_t + 1) * BT, T) - 1 - - if USE_G: - m_t = (i_t * BT + tl.arange(0, BT)) < T - b_g_last = tl.load(g + bos * H + last_idx * H + i_h).to(tl.float32) - 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,)).to(tl.float32) - if USE_EXP2: - b_v = b_v * tl.where(m_t, exp2(b_g_last - b_g), 0)[:, None] - b_g_last = exp2(b_g_last) - else: - b_v = b_v * tl.where(m_t, exp(b_g_last - b_g), 0)[:, None] - b_g_last = exp(b_g_last) - b_h1 *= b_g_last - if K > 64: - b_h2 *= b_g_last - if K > 128: - b_h3 *= b_g_last - if K > 192: - b_h4 *= b_g_last - - 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.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.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.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.0).to( - tl.float32 - ) - if USE_EXP2: - b_h4 *= exp2(b_gk_last4)[:, None] - else: - b_h4 *= exp(b_gk_last4)[:, None] - - b_v = b_v.to(k.dtype.element_ty) - - p_k = tl.make_block_ptr(k, (K, T), (1, stride_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, stride_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, stride_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, stride_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) - - stride_hm_kv = K + V - p_h1 = tl.make_block_ptr(hm, (K, V), (stride_hm_kv, 1), (0, i_v * BLOCK_SIZE), (64, BLOCK_SIZE), (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(hm, (K, V), (stride_hm_kv, 1), (64, i_v * BLOCK_SIZE), (64, BLOCK_SIZE), (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(hm, (K, V), (stride_hm_kv, 1), (128, i_v * BLOCK_SIZE), (64, BLOCK_SIZE), (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(hm, (K, V), (stride_hm_kv, 1), (192, i_v * BLOCK_SIZE), (64, BLOCK_SIZE), (1, 0)) - tl.store(p_h4, b_h4.to(p_h4.dtype.element_ty), boundary_check=(0, 1)) - else: - i_k_col = i_col - tl.cdiv(V, BLOCK_SIZE) - - row = tl.arange(0, BK1) - col = tl.arange(0, BLOCK_SIZE) + i_k_col * BLOCK_SIZE - - b_m = tl.where(row[:, None] == col[None, :], 1.0, 0.0) - - for i_t in range(NT): - p_k = tl.make_block_ptr(k, (T, K), (stride_k, 1), (i_t * BT, 0), (BT, BK1), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - p_w = tl.make_block_ptr(w, (T, K), (stride_k, 1), (i_t * BT, 0), (BT, BK1), (1, 0)) - b_w = tl.load(p_w, boundary_check=(0, 1)) - - last_idx = min((i_t + 1) * BT, T) - 1 - - if USE_G: - m_t = (i_t * BT + tl.arange(0, BT)) < T - b_g_last = tl.load(g + bos * H + last_idx * H + i_h).to(tl.float32) - 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,)).to(tl.float32) - if USE_EXP2: - b_k = b_k * tl.where(m_t, exp2(b_g_last - b_g), 0)[:, None] - b_g_last = exp2(b_g_last) - else: - b_k = b_k * tl.where(m_t, exp(b_g_last - b_g), 0)[:, None] - 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.0).to( - tl.float32 - ) - if USE_EXP2: - b_gk_last = exp2(b_gk_last) - else: - b_gk_last = exp(b_gk_last) - b_diag = tl.where(row[:, None] == row[None, :], b_gk_last[:, None], 0.0) - else: - b_diag = tl.where(row[:, None] == row[None, :], 1.0, 0.0) - - b_kw = tl.dot(tl.trans(b_k.to(b_w.dtype)), b_w) - b_m_i = b_diag - b_kw - b_m = tl.dot(b_m_i.to(tl.float32), b_m.to(tl.float32)) - - stride_hm_kv = K + V - p_m = tl.make_block_ptr(hm + V, (K, K), (stride_hm_kv, 1), (0, i_k_col * BLOCK_SIZE), (BK1, BLOCK_SIZE), (1, 0)) - 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.autotune( - configs=[ - 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"], - use_cuda_graph=USE_CUDA_GRAPH, - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=["pre_or_post_num_ranks", "rank", "NUM_SEQ_ENTRIES"]) -def merge_fwd_bwd_kernel( - h, - ag_hm, - pre_or_post_num_ranks, - rank, - seq_offsets, - init_offsets, - h0_seq_ids, - h0, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BV: tl.constexpr, - BK: tl.constexpr, - FORWARD: tl.constexpr, - INTRACARD_MODE: tl.constexpr, - NUM_SEQ_ENTRIES, - HAS_H0: tl.constexpr, -): - i_v = tl.program_id(0) - if INTRACARD_MODE: - i_seq = tl.program_id(1) - i_h = tl.program_id(2) - - if i_seq >= NUM_SEQ_ENTRIES: - return - - ss_start = tl.load(seq_offsets + i_seq).to(tl.int32) - ss_end = tl.load(seq_offsets + i_seq + 1).to(tl.int32) - init_base = tl.load(init_offsets + i_seq).to(tl.int32) - num_subseqs = ss_end - ss_start - - stride_hm_s = H * K * (V + K) - stride_hm_h = K * (V + K) - - if HAS_H0: - 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), - ) - b_h = tl.load(p_h0, boundary_check=(0, 1)).to(tl.float32) - else: - b_h = tl.zeros([BK, BV], dtype=tl.float32) - - for idx in range(num_subseqs): - i_ss = ss_start + idx - 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), - ) - 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), - ) - 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) - - if idx < num_subseqs - 1: - init_idx = init_base + idx - 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), - ) - tl.store(p_out, b_h.to(p_out.dtype.element_ty), boundary_check=(0, 1)) - else: - i_h = tl.program_id(1) - num_ranks = pre_or_post_num_ranks.to(tl.int32) - h += i_h * K * V - ag_hm += i_h * K * (K + V) - stride = H * K * (K + V) - b_h = tl.zeros([BK, BV], dtype=tl.float32) - for idx in range(num_ranks): - if FORWARD: - cur_rank = rank - num_ranks + idx - else: - cur_rank = rank + num_ranks - idx - p_ag_h = tl.make_block_ptr(ag_hm + cur_rank * stride, (K, V), (K + V, 1), (0, i_v * BV), (BK, BV), (1, 0)) - b_ag_h = tl.load(p_ag_h, boundary_check=(0, 1)) - p_ag_m = tl.make_block_ptr(ag_hm + cur_rank * stride + V, (K, K), (K + V, 1), (0, 0), (BK, BK), (1, 0)) - b_ag_m = tl.load(p_ag_m, boundary_check=(0, 1)) - b_h = tl.dot(b_ag_m.to(tl.float32), b_h.to(tl.float32)) + b_ag_h.to(tl.float32) - p_h = tl.make_block_ptr(h, (K, V), (V, 1), (0, i_v * BV), (BK, BV), (1, 0)) - 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.autotune( - configs=[ - 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] - ], - key=["H", "K", "V", "BT", "BV", "USE_G"], - use_cuda_graph=USE_CUDA_GRAPH, - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=["T"]) -def pre_process_bwd_kernel_stage1( - q, - k, - w, - g, - gk, - do, - dhm, - dv, - cu_seqlens, - scale, - T, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BV: tl.constexpr, - USE_G: tl.constexpr, - USE_GK: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_v, i_h = tl.program_id(0), tl.program_id(1) - i_n = 0 - if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) - T = (eos - bos).to(tl.int32) - NT = tl.cdiv(T, BT) - else: - bos, eos = (i_n * T).to(tl.int64), (i_n * T + T).to(tl.int64) - NT = tl.cdiv(T, BT) - - b_dh1 = tl.zeros([64, BV], dtype=tl.float32) - if K > 64: - b_dh2 = tl.zeros([64, BV], dtype=tl.float32) - if K > 128: - b_dh3 = tl.zeros([64, BV], dtype=tl.float32) - if K > 192: - b_dh4 = tl.zeros([64, BV], dtype=tl.float32) - - q += ((bos * H + i_h) * K).to(tl.int64) - k += ((bos * H + i_h) * K).to(tl.int64) - w += ((bos * H + i_h) * K).to(tl.int64) - do += ((bos * H + i_h) * V).to(tl.int64) - dv += ((bos * H + i_h) * V).to(tl.int64) - dhm += i_h * K * (V + K) - - stride_v = H * V - stride_k = H * K - - for i_t in range(NT - 1, -1, -1): - last_idx = min((i_t + 1) * BT, T) - 1 - if USE_G: - bg_last = tl.load(g + (bos + last_idx) * H + i_h).to(tl.float32) - bg_last_exp = exp(bg_last) - 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,)).to(tl.float32) - b_g_exp = exp(b_g) - - p_dv = tl.make_block_ptr(dv, (T, V), (stride_v, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_do = tl.make_block_ptr(do, (T, V), (stride_v, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - b_do = tl.load(p_do, boundary_check=(0, 1)) - - p_k = tl.make_block_ptr(k, (T, K), (stride_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.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), (stride_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.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), (stride_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.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), (stride_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.0).to(tl.float32) - b_dv += tl.dot(b_k, b_dh4.to(b_k.dtype)) - - if USE_G: - m_t = (i_t * BT + tl.arange(0, BT)) < T - b_dv *= tl.where(m_t, exp(bg_last - b_g), 0)[:, None] - b_dv += tl.load(p_dv, boundary_check=(0, 1)) - - p_w = tl.make_block_ptr(w, (K, T), (1, stride_k), (0, i_t * BT), (64, BT), (0, 1)) - p_q = tl.make_block_ptr(q, (K, T), (1, stride_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: - b_dh1 *= bg_last_exp - b_q = b_q * b_g_exp[None, :] - if USE_GK: - 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, stride_k), (64, i_t * BT), (64, BT), (0, 1)) - p_w = tl.make_block_ptr(w, (K, T), (1, stride_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: - b_dh2 *= bg_last_exp - b_q = b_q * b_g_exp[None, :] - if USE_GK: - 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, stride_k), (128, i_t * BT), (64, BT), (0, 1)) - p_w = tl.make_block_ptr(w, (K, T), (1, stride_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: - b_dh3 *= bg_last_exp - b_q = b_q * b_g_exp[None, :] - if USE_GK: - 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, stride_k), (192, i_t * BT), (64, BT), (0, 1)) - p_w = tl.make_block_ptr(w, (K, T), (1, stride_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: - b_dh4 *= bg_last_exp - b_q = b_q * b_g_exp[None, :] - if USE_GK: - b_dh4 *= exp(b_gk_last4[:, None]) - b_dh4 += 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)) - - p_dh1 = tl.make_block_ptr(dhm, (K, V), (V + K, 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(dhm, (K, V), (V + K, 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(dhm, (K, V), (V + K, 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(dhm, (K, V), (V + K, 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)) - - -@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]) - ], - key=["H", "K", "V", "BT", "USE_EXP2"], - use_cuda_graph=USE_CUDA_GRAPH, - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=["T"]) -def pre_process_bwd_kernel_merged( - q, - k, - w, - g, - gk, - do, - dhm, - dv, - cu_seqlens, - scale, - T, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BLOCK_SIZE: tl.constexpr, - BK1: tl.constexpr, - USE_G: tl.constexpr, - USE_GK: tl.constexpr, - USE_EXP2: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_col, i_h = tl.program_id(0), tl.program_id(1) - i_n = 0 - if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) - T = (eos - bos).to(tl.int32) - NT = tl.cdiv(T, BT) - else: - bos, eos = (i_n * T).to(tl.int64), (i_n * T + T).to(tl.int64) - NT = tl.cdiv(T, BT) - - is_dh_part = i_col * BLOCK_SIZE < V - - q += ((bos * H + i_h) * K).to(tl.int64) - k += ((bos * H + i_h) * K).to(tl.int64) - w += ((bos * H + i_h) * K).to(tl.int64) - dhm += i_h * K * (V + K) - stride_k = H * K - - if is_dh_part: - do += ((bos * H + i_h) * V).to(tl.int64) - dv += ((bos * H + i_h) * V).to(tl.int64) - stride_v = H * V - i_v = i_col - - b_dh1 = tl.zeros([64, BLOCK_SIZE], dtype=tl.float32) - if K > 64: - b_dh2 = tl.zeros([64, BLOCK_SIZE], dtype=tl.float32) - if K > 128: - b_dh3 = tl.zeros([64, BLOCK_SIZE], dtype=tl.float32) - if K > 192: - b_dh4 = tl.zeros([64, BLOCK_SIZE], dtype=tl.float32) - - for i_t in range(NT - 1, -1, -1): - last_idx = min((i_t + 1) * BT, T) - 1 - - if USE_G: - bg_last = tl.load(g + (bos + last_idx) * H + i_h).to(tl.float32) - bg_last_exp = exp(bg_last) - 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,)).to(tl.float32) - b_g_exp = exp(b_g) - - p_dv = tl.make_block_ptr(dv, (T, V), (stride_v, 1), (i_t * BT, i_v * BLOCK_SIZE), (BT, BLOCK_SIZE), (1, 0)) - p_do = tl.make_block_ptr(do, (T, V), (stride_v, 1), (i_t * BT, i_v * BLOCK_SIZE), (BT, BLOCK_SIZE), (1, 0)) - b_do = tl.load(p_do, boundary_check=(0, 1)) - - p_k = tl.make_block_ptr(k, (T, K), (stride_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.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), (stride_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.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), (stride_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.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), (stride_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.0).to(tl.float32) - b_dv += tl.dot(b_k, b_dh4.to(b_k.dtype)) - - if USE_G: - m_t = (i_t * BT + tl.arange(0, BT)) < T - b_dv *= tl.where(m_t, exp(bg_last - b_g), 0)[:, None] - b_dv += tl.load(p_dv, boundary_check=(0, 1)) - - p_w = tl.make_block_ptr(w, (K, T), (1, stride_k), (0, i_t * BT), (64, BT), (0, 1)) - p_q = tl.make_block_ptr(q, (K, T), (1, stride_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: - b_dh1 *= bg_last_exp - b_q = b_q * b_g_exp[None, :] - if USE_GK: - if USE_EXP2: - b_dh1 *= exp2(b_gk_last1[:, None]) - else: - 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, stride_k), (64, i_t * BT), (64, BT), (0, 1)) - p_w = tl.make_block_ptr(w, (K, T), (1, stride_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: - b_dh2 *= bg_last_exp - b_q = b_q * b_g_exp[None, :] - if USE_GK: - if USE_EXP2: - b_dh2 *= exp2(b_gk_last2[:, None]) - else: - 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, stride_k), (128, i_t * BT), (64, BT), (0, 1)) - p_w = tl.make_block_ptr(w, (K, T), (1, stride_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: - b_dh3 *= bg_last_exp - b_q = b_q * b_g_exp[None, :] - if USE_GK: - if USE_EXP2: - b_dh3 *= exp2(b_gk_last3[:, None]) - else: - 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, stride_k), (192, i_t * BT), (64, BT), (0, 1)) - p_w = tl.make_block_ptr(w, (K, T), (1, stride_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: - b_dh4 *= bg_last_exp - b_q = b_q * b_g_exp[None, :] - if USE_GK: - if USE_EXP2: - b_dh4 *= exp2(b_gk_last4[:, None]) - else: - b_dh4 *= exp(b_gk_last4[:, None]) - b_dh4 += 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)) - - p_dh1 = tl.make_block_ptr(dhm, (K, V), (V + K, 1), (0, i_v * BLOCK_SIZE), (64, BLOCK_SIZE), (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(dhm, (K, V), (V + K, 1), (64, i_v * BLOCK_SIZE), (64, BLOCK_SIZE), (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(dhm, (K, V), (V + K, 1), (128, i_v * BLOCK_SIZE), (64, BLOCK_SIZE), (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(dhm, (K, V), (V + K, 1), (192, i_v * BLOCK_SIZE), (64, BLOCK_SIZE), (1, 0)) - tl.store(p_dh4, b_dh4.to(p_dh4.dtype.element_ty), boundary_check=(0, 1)) - else: - i_k_col = i_col - tl.cdiv(V, BLOCK_SIZE) - - row = tl.arange(0, BK1) - col = tl.arange(0, BLOCK_SIZE) + i_k_col * BLOCK_SIZE - - b_m = tl.where(row[:, None] == col[None, :], 1.0, 0.0) - - for _i_t in range(NT): - i_t = NT - 1 - _i_t - - p_k = tl.make_block_ptr(k, (T, K), (stride_k, 1), (i_t * BT, 0), (BT, BK1), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - p_w = tl.make_block_ptr(w, (T, K), (stride_k, 1), (i_t * BT, 0), (BT, BK1), (1, 0)) - b_w = tl.load(p_w, boundary_check=(0, 1)) - - last_idx = min((i_t + 1) * BT, T) - 1 - - if USE_G: - m_t = (i_t * BT + tl.arange(0, BT)) < T - b_g_last = tl.load(g + bos * H + last_idx * H + i_h).to(tl.float32) - 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,)).to(tl.float32) - if USE_EXP2: - b_k = b_k * tl.where(m_t, exp2(b_g_last - b_g), 0)[:, None] - b_g_last = exp2(b_g_last) - else: - b_k = b_k * tl.where(m_t, exp(b_g_last - b_g), 0)[:, None] - 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.0).to( - tl.float32 - ) - if USE_EXP2: - b_gk_last = exp2(b_gk_last) - else: - b_gk_last = exp(b_gk_last) - b_diag = tl.where(row[:, None] == row[None, :], b_gk_last[:, None], 0.0) - else: - b_diag = tl.where(row[:, None] == row[None, :], 1.0, 0.0) - - b_kw = tl.dot(tl.trans(b_w), b_k.to(b_w.dtype)) - b_m_i = b_diag - b_kw - b_m = tl.dot(b_m_i.to(tl.float32), b_m.to(tl.float32)) - - p_m = tl.make_block_ptr(dhm + V, (K, K), (V + K, 1), (0, i_k_col * BLOCK_SIZE), (BK1, BLOCK_SIZE), (1, 0)) - tl.store(p_m, b_m.to(p_m.dtype.element_ty), boundary_check=(0, 1)) - - -def chunk_gated_delta_rule_fwd_h_pre_process( - k: torch.Tensor, - w: torch.Tensor, - u: torch.Tensor, - g: torch.Tensor | None = None, - gk: torch.Tensor | None = None, - chunk_size: int = 64, - cu_seqlens: torch.LongTensor | None = None, - use_exp2: bool = False, - initial_state: torch.Tensor | None = None, - context: FLACPContext = None, -) -> torch.Tensor | None: - if context is None or context.group is None: - return initial_state - assert initial_state is None, "When enable CP, the provided initial_state must be None." - rank = dist.get_rank(group=context.group) - - B, T, H, K, V = *k.shape, u.shape[-1] - BT = chunk_size - BK = triton.next_power_of_2(K) - - if cu_seqlens is None: - N = B - else: - N = len(cu_seqlens) - 1 - assert K <= 256, "current kernel does not support head dimension larger than 256." - - hm = k.new_zeros(H, K, (V + K), dtype=torch.float32) - initial_state = k.new_zeros(N, H, K, V, dtype=torch.float32) - if not context.is_last_rank: - BLOCK_SIZE = 32 if K <= 64 else 64 - grid = (triton.cdiv(V, BLOCK_SIZE) + triton.cdiv(K, BLOCK_SIZE), H) - pre_process_fwd_kernel_merged[grid]( - k=k, - v=u, - w=w, - g=g, - gk=gk, - hm=hm, - cu_seqlens=cu_seqlens[-2:], - T=T, - H=H, - K=K, - V=V, - BT=BT, - BK1=BK, - USE_EXP2=use_exp2, - BLOCK_SIZE=BLOCK_SIZE, - MULTI_SEQS=False, - ) - 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) - - merge_fwd_bwd_kernel[grid]( - h=initial_state[0], - ag_hm=ag_hm, - pre_or_post_num_ranks=context.pre_num_ranks, - rank=rank, - seq_offsets=None, - init_offsets=None, - h0_seq_ids=None, - h0=None, - H=H, - K=K, - V=V, - BK=BK, - FORWARD=True, - INTRACARD_MODE=False, - NUM_SEQ_ENTRIES=0, - ) - return initial_state - - -def chunk_gated_delta_rule_bwd_dhu_pre_process( - q: torch.Tensor, - k: torch.Tensor, - w: torch.Tensor, - do: torch.Tensor, - dv: torch.Tensor, - g: torch.Tensor | None = None, - gk: torch.Tensor | None = None, - scale: float | None = None, - cu_seqlens: torch.LongTensor | None = None, - use_exp2: bool = False, - dht: torch.Tensor | None = None, - initial_state: torch.Tensor | None = None, - context: FLACPContext | None = None, -) -> tuple[torch.Tensor | None, torch.Tensor | None]: - if context is None or context.group is None: - return dht, initial_state - assert dht is None, "When enable CP, the provided dht must be None." - rank = dist.get_rank(context.group) - - B, T, H, K, V = *q.shape, do.shape[-1] - BT = 64 - assert K <= 256, "current kernel does not support head dimension being larger than 256." - BK = triton.next_power_of_2(K) - - if cu_seqlens is None: - N = B - else: - N = len(cu_seqlens) - 1 - - dhm = q.new_zeros(H, K, V + K, dtype=torch.float32) - dht = q.new_zeros(N, H, K, V, dtype=torch.float32) - - if not context.is_first_rank: - BLOCK_SIZE = 32 if K <= 64 else 64 - grid = (triton.cdiv(V, BLOCK_SIZE) + triton.cdiv(K, BLOCK_SIZE), H) - pre_process_bwd_kernel_merged[grid]( - q=q, - k=k, - w=w, - g=g, - gk=gk, - do=do, - dhm=dhm, - dv=dv, - cu_seqlens=cu_seqlens[:2], - scale=scale, - T=T, - H=H, - K=K, - V=V, - BT=BT, - BK1=BK, - USE_EXP2=use_exp2, - BLOCK_SIZE=BLOCK_SIZE, - ) - - 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) - - merge_fwd_bwd_kernel[grid]( - h=dht[-1], - ag_hm=ag_dhm, - pre_or_post_num_ranks=context.post_num_ranks, - rank=rank, - seq_offsets=None, - init_offsets=None, - h0_seq_ids=None, - h0=None, - H=H, - K=K, - V=V, - BK=BK, - FORWARD=False, - INTRACARD_MODE=False, - NUM_SEQ_ENTRIES=0, - ) - - return dht, None - - -def compress_h0(h0: torch.Tensor | None, context: FLACPContext) -> torch.Tensor | None: - if h0 is None or len(context.cu_seqlens) == 2: - return h0 - return h0[:1].clone() - - -def expand_h0(h0: torch.Tensor | None, context: FLACPContext) -> torch.Tensor | None: - if h0 is None or len(context.cu_seqlens) == 2: - return h0 - B = len(context.cu_seqlens) - 1 - expanded_h0 = h0.new_zeros(B, *h0.shape[1:]) - expanded_h0[:1] = h0 - return expanded_h0 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 deleted file mode 100644 index d453f22e..00000000 --- a/src/xorl/ops/linear_attention/ops/gated_delta_rule/__init__.py +++ /dev/null @@ -1,8 +0,0 @@ -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 deleted file mode 100644 index 747758d7..00000000 --- a/src/xorl/ops/linear_attention/ops/gated_delta_rule/chunk.py +++ /dev/null @@ -1,383 +0,0 @@ -# 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 warnings - -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_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 -from xorl.ops.linear_attention.ops.cp.chunk_delta_h import ( - chunk_gated_delta_rule_bwd_dhu_pre_process, - chunk_gated_delta_rule_fwd_h_pre_process, - compress_h0, - expand_h0, -) -from xorl.ops.linear_attention.ops.gated_delta_rule.wy_fast import prepare_wy_repr_bwd, recompute_w_u_fwd -from xorl.ops.linear_attention.ops.utils import chunk_local_cumsum, solve_tril -from xorl.ops.linear_attention.utils import autocast_custom_bwd, autocast_custom_fwd, input_guard - - -def chunk_gated_delta_rule_fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - beta: torch.Tensor, - scale: float, - initial_state: torch.Tensor, - output_final_state: bool, - cu_seqlens: torch.LongTensor | None = None, - cp_context: FLACPContext | None = None, -): - g = chunk_local_cumsum(g, chunk_size=64, cu_seqlens=cu_seqlens) - # obtain WY representation. u is actually the new v. - A = chunk_scaled_dot_kkt_fwd( - k=k, - g=g, - beta=beta, - cu_seqlens=cu_seqlens, - output_dtype=torch.float32, - ) - A = solve_tril( - A=A, - cu_seqlens=cu_seqlens, - output_dtype=k.dtype, - ) - w, u = recompute_w_u_fwd( - k=k, - v=v, - beta=beta, - A=A, - g=g, - cu_seqlens=cu_seqlens, - ) - - if cp_context is not None: - initial_state = chunk_gated_delta_rule_fwd_h_pre_process( - k=k, - w=w, - u=u, - g=g, - cu_seqlens=cu_seqlens, - initial_state=initial_state, - context=cp_context, - ) - - h, v_new, final_state = chunk_gated_delta_rule_fwd_h( - k=k, - w=w, - u=u, - g=g, - initial_state=initial_state, - output_final_state=output_final_state, - cu_seqlens=cu_seqlens, - ) - - if cp_context is not None: - initial_state = compress_h0(initial_state, context=cp_context) - - o = chunk_fwd_o( - q=q, - k=k, - v=v_new, - h=h, - g=g, - scale=scale, - cu_seqlens=cu_seqlens, - ) - return g, o, A, final_state, initial_state - - -def chunk_gated_delta_rule_bwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - beta: torch.Tensor, - A: torch.Tensor, - scale: float, - initial_state: torch.Tensor, - do: torch.Tensor, - dht: torch.Tensor, - cu_seqlens: torch.LongTensor | None = None, - cp_context: FLACPContext | None = None, -): - w, u = recompute_w_u_fwd( - k=k, - v=v, - beta=beta, - A=A, - g=g, - cu_seqlens=cu_seqlens, - ) - - if cp_context is not None: - initial_state = expand_h0(initial_state, context=cp_context) - - h, v_new, _ = chunk_gated_delta_rule_fwd_h( - k=k, - w=w, - u=u, - g=g, - initial_state=initial_state, - output_final_state=False, - cu_seqlens=cu_seqlens, - ) - dv = chunk_bwd_dv_local( - q=q, - k=k, - g=g, - do=do, - scale=scale, - cu_seqlens=cu_seqlens, - ) - - if cp_context is not None: - dht, initial_state = chunk_gated_delta_rule_bwd_dhu_pre_process( - q=q, - k=k, - w=w, - do=do, - dv=dv, - g=g, - scale=scale, - cu_seqlens=cu_seqlens, - dht=dht, - initial_state=initial_state, - context=cp_context, - ) - - dh, dh0, dv = chunk_gated_delta_rule_bwd_dhu( - q=q, - k=k, - w=w, - g=g, - h0=initial_state, - dht=dht, - do=do, - dv=dv, - scale=scale, - cu_seqlens=cu_seqlens, - ) - dq, dk, dw, dg = chunk_bwd_dqkwg( - q=q, - k=k, - v=v_new, - w=w, - g=g, - h=h, - dv=dv, - do=do, - dh=dh, - scale=scale, - cu_seqlens=cu_seqlens, - ) - dk2, dv, db, dg2 = prepare_wy_repr_bwd( - k=k, - v=v, - beta=beta, - g=g, - A=A, - dw=dw, - du=dv, - cu_seqlens=cu_seqlens, - ) - dk.add_(dk2) - dg.add_(dg2) - dg = chunk_local_cumsum(dg, chunk_size=64, reverse=True, cu_seqlens=cu_seqlens) - return dq, dk, dv, db, dg, dh0 - - -class ChunkGatedDeltaRuleFunction(torch.autograd.Function): - @staticmethod - @input_guard - @autocast_custom_fwd - def forward( - ctx, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - beta: torch.Tensor, - scale: float, - initial_state: torch.Tensor, - output_final_state: bool, - cu_seqlens: torch.LongTensor | None = None, - use_qk_l2norm_in_kernel: bool = False, - cp_context: FLACPContext | None = None, - ): - q_rstd, k_rstd = None, None - if use_qk_l2norm_in_kernel: - q, q_rstd = l2norm_fwd(q) - k, k_rstd = l2norm_fwd(k) - - g, o, A, final_state, initial_state = chunk_gated_delta_rule_fwd( - q=q, - k=k, - v=v, - g=g, - beta=beta, - scale=scale, - initial_state=initial_state, - output_final_state=output_final_state, - cu_seqlens=cu_seqlens, - cp_context=cp_context, - ) - ctx.save_for_backward(q, q_rstd, k, k_rstd, v, g, beta, A, initial_state, cu_seqlens) - ctx.scale = scale - ctx.use_qk_l2norm_in_kernel = use_qk_l2norm_in_kernel - ctx.cp_context = cp_context.copy_for_backward() if cp_context is not None else None - return o.to(q.dtype), final_state - - @staticmethod - @input_guard - @autocast_custom_bwd - def backward( - ctx, - do: torch.Tensor, - dht: torch.Tensor, - ): - q, q_rstd, k, k_rstd, v, g, beta, A, initial_state, cu_seqlens = ctx.saved_tensors - dq, dk, dv, db, dg, dh0 = chunk_gated_delta_rule_bwd( - q=q, - k=k, - v=v, - g=g, - beta=beta, - A=A, - scale=ctx.scale, - initial_state=initial_state, - do=do, - dht=dht, - cu_seqlens=cu_seqlens, - cp_context=ctx.cp_context, - ) - if ctx.use_qk_l2norm_in_kernel: - dq = l2norm_bwd(q, q_rstd, dq) - dk = l2norm_bwd(k, k_rstd, dk) - return dq.to(q), dk.to(k), dv.to(v), dg.to(g), db.to(beta), None, dh0, None, None, None, None - - -@torch.compiler.disable -def chunk_gated_delta_rule( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - beta: torch.Tensor, - scale: float = None, - initial_state: torch.Tensor = None, - output_final_state: bool = False, - use_qk_l2norm_in_kernel: bool = False, - cu_seqlens: torch.LongTensor | None = None, - cp_context: FLACPContext | None = None, - **kwargs, -): - r""" - Args: - q (torch.Tensor): - queries of shape `[B, T, H, K]`. - k (torch.Tensor): - keys of shape `[B, T, H, K]`. - v (torch.Tensor): - values of shape `[B, T, H, V]`. - g (torch.Tensor): - (forget) gating tensor (in log space!) of shape `[B, T, H]`. - beta (torch.Tensor): - betas of shape `[B, T, H]`. - scale (Optional[float]): - Scale factor for the RetNet attention scores. - If not provided, it will default to `1 / sqrt(K)`. Default: `None`. - initial_state (Optional[torch.Tensor]): - Initial state of shape `[N, H, K, V]` for `N` input sequences. - For equal-length input sequences, `N` equals the batch size `B`. - Default: `None`. - output_final_state (Optional[bool]): - Whether to output the final state of shape `[N, H, K, V]`. Default: `False`. - use_qk_l2norm_in_kernel (bool): - Whether to apply L2norm to the q/k tensor internally. Default: `False`. - cu_seqlens (torch.LongTensor): - Cumulative sequence lengths of shape `[N+1]` used for variable-length training, - consistent with the FlashAttention API. - - Returns: - o (torch.Tensor): - Outputs of shape `[B, T, H, V]`. - final_state (torch.Tensor): - Final state of shape `[N, H, K, V]` if `output_final_state=True` else `None`. - - Examples:: - >>> import torch - >>> import torch.nn.functional as F - >>> from einops import rearrange - >>> from fla.ops.gated_delta_rule import chunk_gated_delta_rule - # inputs with equal lengths - >>> B, T, H, K, V = 4, 2048, 4, 512, 512 - >>> q = torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda') - >>> k = F.normalize(torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda'), p=2, dim=-1) - >>> v = torch.randn(B, T, H, V, dtype=torch.bfloat16, device='cuda') - >>> beta = torch.rand(B, T, H, dtype=torch.bfloat16, device='cuda').sigmoid() - >>> g = F.logsigmoid(torch.rand(B, T, H, dtype=torch.bfloat16, device='cuda')) - >>> h0 = torch.randn(B, H, K, V, dtype=torch.bfloat16, device='cuda') - >>> o, ht = chunk_gated_delta_rule( - q, k, v, g, beta, - initial_state=h0, - output_final_state=True - ) - # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required - >>> q, k, v, beta, g = map(lambda x: rearrange(x, 'b t ... -> 1 (b t) ...'), (q, k, v, beta, g)) - # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected - >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long) - >>> o, ht = chunk_gated_delta_rule( - q, k, v, g, beta, - initial_state=h0, - output_final_state=True, - cu_seqlens=cu_seqlens - ) - """ - 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.", - ) - - if cp_context is not None: - assert initial_state is None, "Initial state is not supported for CP" - assert output_final_state is False, "Output final state is not supported for CP" - assert cp_context.cu_seqlens is not None, "cu_seqlens is required for CP" - cu_seqlens = cp_context.cu_seqlens - - if cu_seqlens is not None: - if q.shape[0] != 1: - raise ValueError( - f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`." - f"Please flatten variable-length inputs before processing.", - ) - if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: - raise ValueError( - f"The number of initial states is expected to be equal to the number of input sequences, " - f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.", - ) - if scale is None: - scale = k.shape[-1] ** -0.5 - o, final_state = ChunkGatedDeltaRuleFunction.apply( - q, - k, - v, - g, - beta, - scale, - initial_state, - output_final_state, - cu_seqlens, - use_qk_l2norm_in_kernel, - cp_context, - ) - return o, final_state 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 deleted file mode 100644 index 043209a2..00000000 --- a/src/xorl/ops/linear_attention/ops/gated_delta_rule/fused_recurrent.py +++ /dev/null @@ -1,249 +0,0 @@ -from __future__ import annotations - -# 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 - -from xorl.ops.linear_attention.ops.utils.op import exp -from xorl.ops.linear_attention.utils import input_guard - - -@triton.heuristics( - { - "USE_G": lambda args: args["g"] is not None, - "USE_GK": lambda args: args["gk"] is not None, - "USE_GV": lambda args: args["gv"] is not None, - "USE_INITIAL_STATE": lambda args: args["h0"] is not None, - "STORE_FINAL_STATE": lambda args: args["ht"] is not None, - "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, - } -) -@triton.jit(do_not_specialize=["T"]) -def fused_recurrent_gated_delta_rule_fwd_kernel( - q, - k, - v, - g, - gk, - gv, - beta, - o, - h0, - ht, - cu_seqlens, - scale, - T, - H: tl.constexpr, - HV: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - USE_G: tl.constexpr, - USE_GK: tl.constexpr, - USE_GV: tl.constexpr, - USE_QK_L2NORM_IN_KERNEL: tl.constexpr, - IS_BETA_HEADWISE: tl.constexpr, - USE_INITIAL_STATE: tl.constexpr, - STORE_FINAL_STATE: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_v, i_nh = tl.program_id(0), tl.program_id(1) - i_n, i_hv = i_nh // HV, i_nh % HV - i_h = i_hv // (HV // H) - - if IS_VARLEN: - bos = tl.load(cu_seqlens + i_n).to(tl.int64) - eos = tl.load(cu_seqlens + i_n + 1).to(tl.int64) - T = eos - bos - else: - bos, eos = i_n * T, i_n * T + T - - o_k = tl.arange(0, BK) - o_v = i_v * BV + tl.arange(0, BV) - - p_q = q + (bos * H + i_h) * K + o_k - p_k = k + (bos * H + i_h) * K + o_k - p_v = v + (bos * HV + i_hv) * V + o_v - if USE_G: - p_g = g + bos * HV + i_hv - if USE_GK: - p_gk = gk + (bos * HV + i_hv) * K + o_k - if USE_GV: - p_gv = gv + (bos * HV + i_hv) * V + o_v - if IS_BETA_HEADWISE: - p_beta = beta + bos * HV + i_hv - else: - p_beta = beta + (bos * HV + i_hv) * V + o_v - p_o = o + (bos * HV + i_hv) * V + o_v - - mask_k = o_k < K - mask_v = o_v < V - mask_h = mask_k[:, None] & mask_v[None, :] - - b_h = tl.zeros([BK, BV], dtype=tl.float32) - if USE_INITIAL_STATE: - p_h0 = h0 + i_nh * K * V + o_k[:, None] * V + o_v[None, :] - b_h += tl.load(p_h0, mask=mask_h, other=0).to(tl.float32) - - for _ in tl.range(0, T): - b_q = tl.load(p_q, mask=mask_k, other=0).to(tl.float32) - b_k = tl.load(p_k, mask=mask_k, other=0).to(tl.float32) - b_v = tl.load(p_v, mask=mask_v, other=0).to(tl.float32) - if USE_QK_L2NORM_IN_KERNEL: - b_q = b_q / tl.sqrt(tl.sum(b_q * b_q) + 1e-6) - b_k = b_k / tl.sqrt(tl.sum(b_k * b_k) + 1e-6) - b_q = b_q * scale - - if IS_BETA_HEADWISE: - b_beta = tl.load(p_beta).to(tl.float32) - else: - b_beta = tl.load(p_beta, mask=mask_v, other=0).to(tl.float32) - - if USE_G: - b_h *= exp(tl.load(p_g).to(tl.float32)) - if USE_GK: - b_h *= exp(tl.load(p_gk).to(tl.float32)[:, None]) - if USE_GV: - b_h *= exp(tl.load(p_gv).to(tl.float32)[None, :]) - - b_v = b_beta * (b_v - tl.sum(b_h * b_k[:, None], 0)) - b_h += b_k[:, None] * b_v - b_o = tl.sum(b_h * b_q[:, None], 0) - tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=mask_v) - - p_q += H * K - p_k += H * K - p_v += HV * V - if USE_G: - p_g += HV - if USE_GK: - p_gk += HV * K - if USE_GV: - p_gv += HV * V - p_beta += HV * (1 if IS_BETA_HEADWISE else V) - p_o += HV * V - - if STORE_FINAL_STATE: - p_ht = ht + i_nh * K * V + o_k[:, None] * V + o_v[None, :] - tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h) - - -def fused_recurrent_gated_delta_rule_fwd( - 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, -) -> tuple[torch.Tensor, torch.Tensor | None]: - B, T, H, K, V = *k.shape, v.shape[-1] - HV = v.shape[2] - N = B if cu_seqlens is None else len(cu_seqlens) - 1 - BK = triton.next_power_of_2(K) - BV = min(8, triton.next_power_of_2(V)) if gv is None else triton.next_power_of_2(V) - NV = triton.cdiv(V, BV) - - o = torch.empty_like(v) - final_state = q.new_empty(N, HV, K, V, dtype=torch.float32) if output_final_state else None - - grid = (NV, N * HV) - fused_recurrent_gated_delta_rule_fwd_kernel[grid]( - q=q, - k=k, - v=v, - g=g, - gk=gk, - gv=gv, - beta=beta, - o=o, - h0=initial_state, - ht=final_state, - cu_seqlens=cu_seqlens, - scale=scale, - T=T, - H=H, - HV=HV, - K=K, - V=V, - BK=BK, - BV=BV, - IS_BETA_HEADWISE=beta.ndim != v.ndim, - USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel, - num_warps=1, - num_stages=3, - ) - return o, final_state - - -class FusedRecurrentFunction(torch.autograd.Function): - @staticmethod - @input_guard - def forward( - ctx, - 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, - ): - del ctx - return fused_recurrent_gated_delta_rule_fwd( - q=q, - k=k, - v=v, - g=g, - gk=gk, - gv=gv, - beta=beta, - scale=scale, - initial_state=initial_state, - output_final_state=output_final_state, - use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, - cu_seqlens=cu_seqlens, - ) - - -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 deleted file mode 100644 index 18424ecc..00000000 --- a/src/xorl/ops/linear_attention/ops/gated_delta_rule/wy_fast.py +++ /dev/null @@ -1,345 +0,0 @@ -# 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 triton.language as tl - -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, -) - - -if IS_NVIDIA_BLACKWELL: - """ - Compute tl.dot with SM100 workaround. - - On SM100 (Blackwell) GPUs, wraps the result in inline assembly to prevent - the TritonGPUHoistTMEMAlloc pass from incorrectly fusing add and dot operations. - See: https://github.com/fla-org/flash-linear-attention/issues/638 - - 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( - asm="mov.f32 $0, $1;", - constraints="=r,r", - args=[tl.dot(a, b)], - dtype=tl.float32, - is_pure=True, - pack=1, - ) -else: - - @triton.jit - def safe_dot(a, b): - return tl.dot(a, b) - - -# SM90-class GPUs can hit illegal memory access in GDN backward autotuning -# with the 4-warps/2-stages variant on real Qwen3.5 shapes (e.g. T=4096, H=4, D=128). -# Keep the bad candidate out until the Triton codegen issue is resolved upstream. -PREPARE_WY_REPR_BWD_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] - if not (IS_NVIDIA_HOPPER and num_warps == 4 and num_stages == 2) -] - - -@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"], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=["T"]) -def recompute_w_u_fwd_kernel( - k, - v, - beta, - w, - u, - A, - g, - cu_seqlens, - chunk_indices, - T, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - USE_G: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_t, i_bh = tl.program_id(0), tl.program_id(1) - i_b, i_h = i_bh // H, i_bh % H - if IS_VARLEN: - i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - 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,)) - 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)) - 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)) - 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,)) - 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)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_kb = b_k * b_b[:, None] - if USE_G: - b_kb *= b_g[:, None] - b_w = tl.dot(b_A, b_kb.to(b_k.dtype)) - 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.autotune( - configs=PREPARE_WY_REPR_BWD_AUTOTUNE_CONFIGS, - key=["H", "K", "V", "BT", "BK", "BV", "IS_VARLEN"], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=["T"]) -def prepare_wy_repr_bwd_kernel( - k, - v, - beta, - g, - A, - dw, - du, - dk, - dv, - db, - dg, - cu_seqlens, - chunk_indices, - T, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - USE_G: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_t, i_bh = tl.program_id(0), tl.program_id(1) - i_b, i_h = i_bh // H, i_bh % H - if IS_VARLEN: - i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - 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_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) - b_A = tl.load(p_A, boundary_check=(0, 1)) - 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,)) - 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)) - # [BT, BK] - b_k = tl.load(p_k, boundary_check=(0, 1)) - if USE_G: - b_kbg = b_k * (b_b * b_g_exp)[:, None] - else: - b_kbg = b_k * b_b[:, None] - b_dw = tl.load(p_dw, boundary_check=(0, 1)) - - b_dA += tl.dot(b_dw, tl.trans(b_kbg).to(b_dw.dtype)) - b_dkbg = tl.dot(b_A, b_dw) - if USE_G: - b_dk = b_dkbg * (b_g_exp * b_b)[:, None] - b_db += tl.sum(b_dkbg * b_k * b_g_exp[:, None], 1) - b_dg += tl.sum(b_dkbg * b_kbg, 1) - 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)) - 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)) - b_dA += tl.dot(b_du, tl.trans(b_vb)) - b_dvb = tl.dot(b_A, b_du) - b_dv = b_dvb * b_b[:, None] - b_db += tl.sum(b_dvb * b_v, 1) - tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) - - o_t = i_t * BT + tl.arange(0, BT) - m_t = o_t < T - m_A = (o_t[:, None] > o_t[None, :]) & (m_t[:, None] & m_t) - b_dA = tl.where(m_A, b_dA, 0) - b_dA = tl.dot(b_dA.to(b_A.dtype), b_A) - b_dA = tl.dot(b_A, b_dA.to(b_A.dtype)) - - if USE_G: - b_dA *= exp(b_g[:, None] - b_g[None, :]) - - b_A = tl.zeros([BT, BT], dtype=tl.float32) - b_dA = tl.where(m_A, -b_dA, 0).to(k.dtype.element_ty) - - 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)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_kt = tl.trans(b_k) - b_ktb = b_kt * b_b[None, :] - - b_A += tl.dot(b_k, b_kt) - b_dkb = tl.dot(b_dA, b_k) - b_db += tl.sum(b_dkb * b_k, 1) - b_dk = b_dkb * b_b[:, None] + tl.trans(tl.dot(b_ktb.to(b_dA.dtype), b_dA)) - b_dk += tl.load(p_dk, boundary_check=(0, 1)) - - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_db, b_db.to(p_db.dtype.element_ty), boundary_check=(0,)) - - 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,)) - 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,)) - - -def recompute_w_u_fwd( - k: torch.Tensor, - v: torch.Tensor, - beta: torch.Tensor, - A: torch.Tensor, - g: torch.Tensor | None = None, - cu_seqlens: torch.LongTensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor]: - B, T, H, K, V = *k.shape, v.shape[-1] - BT = A.shape[-1] - BK = 64 - BV = 64 - - 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) - - w = torch.empty_like(k) - u = torch.empty_like(v) - recompute_w_u_fwd_kernel[(NT, B * H)]( - k=k, - v=v, - beta=beta, - w=w, - u=u, - A=A, - g=g, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - T=T, - H=H, - K=K, - V=V, - BT=BT, - BK=BK, - BV=BV, - ) - return w, u - - -def prepare_wy_repr_bwd( - k: torch.Tensor, - v: torch.Tensor, - beta: torch.Tensor, - A: torch.Tensor, - dw: torch.Tensor, - du: torch.Tensor, - g: torch.Tensor = None, - cu_seqlens: torch.LongTensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - B, T, H, K, V = *k.shape, v.shape[-1] - BT = 64 - 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) - CONST_TILING = 64 if check_shared_mem() else 32 - BK = min(max(triton.next_power_of_2(K), 16), CONST_TILING) - BV = min(max(triton.next_power_of_2(V), 16), CONST_TILING) - - dk = torch.empty_like(k) - dv = torch.empty_like(v) - dg = torch.empty_like(g) if g is not None else None - db = torch.empty_like(beta) - prepare_wy_repr_bwd_kernel[(NT, B * H)]( - k=k, - v=v, - beta=beta, - g=g, - A=A, - dw=dw, - du=du, - dk=dk, - dv=dv, - db=db, - dg=dg, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - T=T, - H=H, - K=K, - V=V, - BT=BT, - BK=BK, - BV=BV, - ) - return dk, dv, db, dg diff --git a/src/xorl/ops/linear_attention/ops/utils/__init__.py b/src/xorl/ops/linear_attention/ops/utils/__init__.py deleted file mode 100644 index e754859e..00000000 --- a/src/xorl/ops/linear_attention/ops/utils/__init__.py +++ /dev/null @@ -1,15 +0,0 @@ -from .cumsum import chunk_local_cumsum -from .index import prepare_chunk_indices, prepare_chunk_offsets -from .op import exp, exp2, make_tensor_descriptor -from .solve_tril import solve_tril - - -__all__ = [ - "chunk_local_cumsum", - "prepare_chunk_indices", - "prepare_chunk_offsets", - "exp", - "exp2", - "make_tensor_descriptor", - "solve_tril", -] diff --git a/src/xorl/ops/linear_attention/ops/utils/cumsum.py b/src/xorl/ops/linear_attention/ops/utils/cumsum.py deleted file mode 100644 index e27658a3..00000000 --- a/src/xorl/ops/linear_attention/ops/utils/cumsum.py +++ /dev/null @@ -1,474 +0,0 @@ -# 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 triton.language as tl - -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.autotune( - 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"]) -def chunk_local_cumsum_scalar_kernel( - s, - o, - scale, - cu_seqlens, - chunk_indices, - T, - B: tl.constexpr, - H: tl.constexpr, - BT: tl.constexpr, - REVERSE: tl.constexpr, - HAS_SCALE: tl.constexpr, - IS_VARLEN: tl.constexpr, - HEAD_FIRST: tl.constexpr, -): - i_t, i_bh = tl.program_id(0), tl.program_id(1) - i_b, i_h = i_bh // H, i_bh % H - if IS_VARLEN: - i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - else: - 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,)) - 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,)) - # [BT] - b_s = tl.load(p_s, boundary_check=(0,)).to(tl.float32) - b_o = tl.cumsum(b_s, axis=0) - if REVERSE: - b_z = tl.sum(b_s, axis=0) - b_o = -b_o + b_z[None] + b_s - if HAS_SCALE: - b_o *= scale - 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.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"], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=["T"]) -def chunk_local_cumsum_vector_kernel( - s, - o, - scale, - cu_seqlens, - chunk_indices, - T, - B: tl.constexpr, - H: tl.constexpr, - S: tl.constexpr, - BT: tl.constexpr, - BS: tl.constexpr, - REVERSE: tl.constexpr, - HAS_SCALE: tl.constexpr, - IS_VARLEN: tl.constexpr, - HEAD_FIRST: tl.constexpr, -): - i_s, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_b, i_h = i_bh // H, i_bh % H - if IS_VARLEN: - i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - else: - 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)) - 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)) - # [BT, BS] - b_s = tl.load(p_s, boundary_check=(0, 1)).to(tl.float32) - if REVERSE: - b_o = tl.cumsum(b_s, axis=0, reverse=True) - else: - b_o = tl.cumsum(b_s, axis=0) - if HAS_SCALE: - b_o *= scale - 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.autotune( - configs=[ - 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"], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=["T"]) -def chunk_global_cumsum_scalar_kernel( - s, - o, - scale, - cu_seqlens, - T, - B: tl.constexpr, - H: tl.constexpr, - BT: tl.constexpr, - REVERSE: tl.constexpr, - HAS_SCALE: tl.constexpr, - IS_VARLEN: tl.constexpr, - HEAD_FIRST: tl.constexpr, -): - i_nh = tl.program_id(0) - i_n, i_h = i_nh // H, i_nh % H - if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - else: - bos, eos = i_n * T, i_n * T + T - T = eos - bos - - b_z = tl.zeros([], dtype=tl.float32) - NT = tl.cdiv(T, BT) - 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,)) - 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,)) - 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) - if REVERSE: - b_o = -b_o + b_ss + b_s - b_o += b_z - if i_c >= 0: - b_z += b_ss - if HAS_SCALE: - b_o *= scale - 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.autotune( - configs=[ - 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"], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=["T"]) -def chunk_global_cumsum_vector_kernel( - s, - o, - scale, - cu_seqlens, - T, - B: tl.constexpr, - H: tl.constexpr, - S: tl.constexpr, - BT: tl.constexpr, - BS: tl.constexpr, - REVERSE: tl.constexpr, - HAS_SCALE: tl.constexpr, - IS_VARLEN: tl.constexpr, - HEAD_FIRST: tl.constexpr, -): - i_s, i_nh = tl.program_id(0), tl.program_id(1) - i_n, i_h = i_nh // H, i_nh % H - if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - else: - bos, eos = i_n * T, i_n * T + T - T = eos - bos - - b_z = tl.zeros([BS], dtype=tl.float32) - NT = tl.cdiv(T, BT) - 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)) - 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)) - # [BT, BS] - b_s = tl.load(p_s, boundary_check=(0, 1)).to(tl.float32) - if REVERSE: - b_c = b_z[None, :] + tl.cumsum(b_s, axis=0, reverse=True) - else: - b_c = b_z[None, :] + tl.cumsum(b_s, axis=0) - if HAS_SCALE: - b_c *= scale - tl.store(p_o, b_c.to(p_o.dtype.element_ty), boundary_check=(0, 1)) - b_z += tl.sum(b_s, 0) - - -def chunk_local_cumsum_scalar( - g: torch.Tensor, - chunk_size: int, - reverse: bool = False, - scale: float = None, - cu_seqlens: torch.Tensor | None = None, - head_first: bool = False, - output_dtype: torch.dtype | None = torch.float, - chunk_indices: torch.LongTensor | None = None, -) -> torch.Tensor: - if head_first: - 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" - BT = chunk_size - 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) - g_org, g = g, torch.empty_like(g, dtype=output_dtype or g.dtype) - grid = (NT, B * H) - chunk_local_cumsum_scalar_kernel[grid]( - s=g_org, - o=g, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - T=T, - B=B, - H=H, - BT=BT, - HEAD_FIRST=head_first, - REVERSE=reverse, - ) - return g - - -def chunk_local_cumsum_vector( - g: torch.Tensor, - chunk_size: int, - reverse: bool = False, - scale: float = None, - cu_seqlens: torch.Tensor | None = None, - head_first: bool = False, - output_dtype: torch.dtype | None = torch.float, - chunk_indices: torch.LongTensor | None = None, -) -> torch.Tensor: - if head_first: - B, H, T, S = g.shape - else: - B, T, H, S = g.shape - BT = chunk_size - 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" - - 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 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]( - s=g_org, - o=g, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - T=T, - B=B, - H=H, - S=S, - BT=BT, - HEAD_FIRST=head_first, - REVERSE=reverse, - ) - return g - - -@input_guard -def chunk_global_cumsum_scalar( - s: torch.Tensor, - reverse: bool = False, - cu_seqlens: torch.Tensor | None = None, - scale: float = None, - head_first: bool = False, - output_dtype: torch.dtype | None = torch.float, -) -> torch.Tensor: - if head_first: - B, H, T = s.shape - else: - B, T, H = s.shape - N = len(cu_seqlens) - 1 if cu_seqlens is not None else B - - z = torch.empty_like(s, dtype=output_dtype or s.dtype) - grid = (N * H,) - chunk_global_cumsum_scalar_kernel[grid]( - s=s, - o=z, - scale=scale, - cu_seqlens=cu_seqlens, - T=T, - B=B, - H=H, - HEAD_FIRST=head_first, - REVERSE=reverse, - ) - return z - - -@input_guard -def chunk_global_cumsum_vector( - s: torch.Tensor, - reverse: bool = False, - cu_seqlens: torch.Tensor | None = None, - scale: float = None, - head_first: bool = False, - output_dtype: torch.dtype | None = torch.float, -) -> torch.Tensor: - if head_first: - B, H, T, S = s.shape - else: - B, T, H, S = s.shape - N = len(cu_seqlens) - 1 if cu_seqlens is not None else B - BS = min(32, triton.next_power_of_2(S)) - - z = torch.empty_like(s, dtype=output_dtype or s.dtype) - grid = (triton.cdiv(S, BS), N * H) - chunk_global_cumsum_vector_kernel[grid]( - s=s, - o=z, - scale=scale, - cu_seqlens=cu_seqlens, - T=T, - B=B, - H=H, - S=S, - BS=BS, - HEAD_FIRST=head_first, - REVERSE=reverse, - ) - return z - - -@input_guard -def chunk_global_cumsum( - s: torch.Tensor, - reverse: bool = False, - cu_seqlens: torch.Tensor | None = None, - scale: float = None, - head_first: bool = False, - output_dtype: torch.dtype | None = torch.float, -) -> torch.Tensor: - if cu_seqlens is not None: - assert s.shape[0] == 1, "Only batch size 1 is supported when cu_seqlens are provided" - if len(s.shape) == 3: - return chunk_global_cumsum_scalar( - s=s, - reverse=reverse, - cu_seqlens=cu_seqlens, - scale=scale, - head_first=head_first, - output_dtype=output_dtype, - ) - elif len(s.shape) == 4: - return chunk_global_cumsum_vector( - s=s, - reverse=reverse, - cu_seqlens=cu_seqlens, - scale=scale, - head_first=head_first, - output_dtype=output_dtype, - ) - else: - raise ValueError( - f"Unsupported input shape {s.shape}, " - f"which should be [B, T, H]/[B, T, H, D] if `head_first=False` " - f"or [B, H, T]/[B, H, T, D] otherwise", - ) - - -@input_guard -def chunk_local_cumsum( - g: torch.Tensor, - chunk_size: int, - reverse: bool = False, - scale: float = None, - cu_seqlens: torch.Tensor | None = None, - head_first: bool = False, - output_dtype: torch.dtype | None = torch.float, - chunk_indices: torch.LongTensor | None = None, - **kwargs, -) -> torch.Tensor: - if cu_seqlens is not None: - assert g.shape[0] == 1, "Only batch size 1 is supported when cu_seqlens are provided" - if len(g.shape) == 3: - return chunk_local_cumsum_scalar( - g=g, - chunk_size=chunk_size, - reverse=reverse, - scale=scale, - cu_seqlens=cu_seqlens, - head_first=head_first, - output_dtype=output_dtype, - chunk_indices=chunk_indices, - ) - elif len(g.shape) == 4: - return chunk_local_cumsum_vector( - g=g, - chunk_size=chunk_size, - reverse=reverse, - scale=scale, - cu_seqlens=cu_seqlens, - head_first=head_first, - output_dtype=output_dtype, - chunk_indices=chunk_indices, - ) - else: - raise ValueError( - f"Unsupported input shape {g.shape}, " - f"which should be (B, T, H, D) if `head_first=False` " - f"or (B, H, T, D) otherwise", - ) diff --git a/src/xorl/ops/linear_attention/ops/utils/index.py b/src/xorl/ops/linear_attention/ops/utils/index.py deleted file mode 100644 index bcae2f2a..00000000 --- a/src/xorl/ops/linear_attention/ops/utils/index.py +++ /dev/null @@ -1,38 +0,0 @@ -# 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 torch.nn.functional as F -import triton - -from xorl.ops.linear_attention.utils import tensor_cache - - -@tensor_cache -def prepare_lens(cu_seqlens: torch.LongTensor) -> torch.LongTensor: - return torch.diff(cu_seqlens) - - -@tensor_cache -def prepare_chunk_indices( - cu_seqlens: torch.LongTensor, - chunk_size: int, - 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() - ] - ) - 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) - - -@tensor_cache -def prepare_chunk_offsets( - cu_seqlens: torch.LongTensor, - chunk_size: int, -) -> torch.LongTensor: - return F.pad(triton.cdiv(prepare_lens(cu_seqlens), chunk_size), (1, 0), value=0).cumsum(-1) diff --git a/src/xorl/ops/linear_attention/ops/utils/op.py b/src/xorl/ops/linear_attention/ops/utils/op.py deleted file mode 100644 index 200575c7..00000000 --- a/src/xorl/ops/linear_attention/ops/utils/op.py +++ /dev/null @@ -1,55 +0,0 @@ -from __future__ import annotations - -# 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 - def exp(x): - return tldevice.fast_expf(x.to(tl.float32)) - - @triton.jit - def exp2(x): - return tldevice.exp2(x.to(tl.float32)) - -else: - - @triton.jit - def exp(x): - return tl.exp(x.to(tl.float32)) - - @triton.jit - def exp2(x): - return tl.math.exp2(x.to(tl.float32)) - - -if not IS_GATHER_SUPPORTED: - - @triton.jit - def gather(src, index, axis, _builder=None): - del src, index, axis, _builder - return None - -else: - gather = tl.gather - - -if hasattr(triton.language, "_experimental_make_tensor_descriptor"): - make_tensor_descriptor = triton.language._experimental_make_tensor_descriptor -elif hasattr(triton.language, "make_tensor_descriptor"): - make_tensor_descriptor = triton.language.make_tensor_descriptor -else: - - @triton.jit - def make_tensor_descriptor(base, shape, strides, block_shape, _builder=None): - del base, shape, strides, block_shape, _builder - return None diff --git a/src/xorl/ops/linear_attention/ops/utils/solve_tril.py b/src/xorl/ops/linear_attention/ops/utils/solve_tril.py deleted file mode 100644 index d1a34f59..00000000 --- a/src/xorl/ops/linear_attention/ops/utils/solve_tril.py +++ /dev/null @@ -1,399 +0,0 @@ -# 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 os - -import torch -import triton -import triton.language as tl - -from xorl.ops.linear_attention.ops.utils.index import prepare_chunk_indices -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"], ( - 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.autotune( - configs=[ - 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"], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=["T"]) -def solve_tril_16x16_kernel( - A, - Ai, - cu_seqlens, - chunk_indices, - T, - H: tl.constexpr, - BT: tl.constexpr, - USE_TMA: tl.constexpr, - IS_VARLEN: tl.constexpr, - DOT_PRECISION: tl.constexpr, -): - i_t, i_bh = tl.program_id(0), tl.program_id(1) - i_b, i_h = i_bh // H, i_bh % H - if IS_VARLEN: - i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - else: - bos, eos = i_b * T, i_b * T + T - o_i = tl.arange(0, 16) - 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 - - 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)) - # [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]) - 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.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)) - 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.autotune( - configs=[ - 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"], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=["T"]) -def merge_16x16_to_32x32_inverse_kernel( - A, - Ai, - cu_seqlens, - chunk_indices, - T, - H: tl.constexpr, - BT: tl.constexpr, - USE_TMA: tl.constexpr, - IS_VARLEN: tl.constexpr, - DOT_PRECISION: tl.constexpr, -): - i_t, i_bh = tl.program_id(0), tl.program_id(1) - i_b, i_h = i_bh // H, i_bh % H - if IS_VARLEN: - i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - else: - bos, eos = i_b * T, i_b * T + T - - o_i = tl.arange(0, 16) - m_A = o_i[:, None] > o_i[None, :] - m_I = o_i[:, None] == o_i[None, :] - A += (bos * H + i_h) * BT - 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)) - 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]) - 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) - - # [16, 16] - b_Ai_11 = -tl.where(m_A, b_Ai_11, 0) - 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.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.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) - - b_Ai_11 += m_I - 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)) - 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) - - 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)) - 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)) - else: - desc_o.store([i_t * BT + 0, 0], b_Ai_11.to(desc_o.dtype, fp_downcast_rounding="rtne")) - desc_o.store([i_t * BT + 16, 0], b_Ai_21.to(desc_o.dtype, fp_downcast_rounding="rtne")) - 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.autotune( - configs=[ - 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"], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=["T"]) -def merge_16x16_to_64x64_inverse_kernel( - A, - Ai, - cu_seqlens, - chunk_indices, - T, - H: tl.constexpr, - BT: tl.constexpr, - USE_TMA: tl.constexpr, - IS_VARLEN: tl.constexpr, - DOT_PRECISION: tl.constexpr, -): - i_t, i_bh = tl.program_id(0), tl.program_id(1) - i_b, i_h = i_bh // H, i_bh % H - if IS_VARLEN: - i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - else: - bos, eos = i_b * T, i_b * T + T - - o_i = tl.arange(0, 16) - m_A = o_i[:, None] > o_i[None, :] - m_I = o_i[:, None] == o_i[None, :] - A += (bos * H + i_h) * BT - 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)) - 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]) - 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) - b_Ai_44 = desc.load([i_t * BT + 48, 48]).to(tl.float32) - - # [16, 16] - b_Ai_11 = -tl.where(m_A, b_Ai_11, 0) - b_Ai_22 = -tl.where(m_A, b_Ai_22, 0) - b_Ai_33 = -tl.where(m_A, b_Ai_33, 0) - 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.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.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.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.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 - b_Ai_22 += m_I - b_Ai_33 += m_I - 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)) - 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) - b_A_41 = tl.load(p_A_41, boundary_check=(0, 1)).to(tl.float32) - b_A_42 = tl.load(p_A_42, boundary_check=(0, 1)).to(tl.float32) - b_A_43 = tl.load(p_A_43, boundary_check=(0, 1)).to(tl.float32) - else: - b_A_21 = desc.load([i_t * BT + 16, 0]).to(tl.float32) - b_A_31 = desc.load([i_t * BT + 32, 0]).to(tl.float32) - b_A_32 = desc.load([i_t * BT + 32, 16]).to(tl.float32) - b_A_41 = desc.load([i_t * BT + 48, 0]).to(tl.float32) - b_A_42 = desc.load([i_t * BT + 48, 16]).to(tl.float32) - b_A_43 = desc.load([i_t * BT + 48, 32]).to(tl.float32) - - b_Ai_21 = -tl.dot(tl.dot(b_Ai_22, b_A_21, input_precision=DOT_PRECISION), b_Ai_11, input_precision=DOT_PRECISION) - b_Ai_32 = -tl.dot(tl.dot(b_Ai_33, b_A_32, input_precision=DOT_PRECISION), b_Ai_22, input_precision=DOT_PRECISION) - b_Ai_43 = -tl.dot(tl.dot(b_Ai_44, b_A_43, input_precision=DOT_PRECISION), b_Ai_33, input_precision=DOT_PRECISION) - - 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), - 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), - 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), - 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)) - 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)) - tl.store(p_Ai_44, b_Ai_44.to(p_Ai_44.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)) - tl.store(p_Ai_31, b_Ai_31.to(p_Ai_31.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) - tl.store(p_Ai_32, b_Ai_32.to(p_Ai_32.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) - tl.store(p_Ai_41, b_Ai_41.to(p_Ai_41.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) - tl.store(p_Ai_42, b_Ai_42.to(p_Ai_42.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) - tl.store(p_Ai_43, b_Ai_43.to(p_Ai_43.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) - else: - desc_o.store([i_t * BT + 0, 0], b_Ai_11.to(desc_o.dtype, fp_downcast_rounding="rtne")) - desc_o.store([i_t * BT + 16, 16], b_Ai_22.to(desc_o.dtype, fp_downcast_rounding="rtne")) - desc_o.store([i_t * BT + 32, 32], b_Ai_33.to(desc_o.dtype, fp_downcast_rounding="rtne")) - desc_o.store([i_t * BT + 48, 48], b_Ai_44.to(desc_o.dtype, fp_downcast_rounding="rtne")) - desc_o.store([i_t * BT + 16, 0], b_Ai_21.to(desc_o.dtype, fp_downcast_rounding="rtne")) - desc_o.store([i_t * BT + 32, 0], b_Ai_31.to(desc_o.dtype, fp_downcast_rounding="rtne")) - desc_o.store([i_t * BT + 32, 16], b_Ai_32.to(desc_o.dtype, fp_downcast_rounding="rtne")) - desc_o.store([i_t * BT + 48, 0], b_Ai_41.to(desc_o.dtype, fp_downcast_rounding="rtne")) - desc_o.store([i_t * BT + 48, 16], b_Ai_42.to(desc_o.dtype, fp_downcast_rounding="rtne")) - desc_o.store([i_t * BT + 48, 32], b_Ai_43.to(desc_o.dtype, fp_downcast_rounding="rtne")) - - -@input_guard -def solve_tril( - A: torch.Tensor, - cu_seqlens: torch.Tensor | None = None, - chunk_indices: torch.LongTensor | None = None, - output_dtype: torch.dtype = torch.float, -) -> torch.Tensor: - """ - Compute the inverse of the matrix I + A - A should be strictly lower triangular, i.e., A.triu() == 0. - - Args: - A (torch.Tensor): - [B, T, H, BT], where BT should only be 16, 32, or 64. - cu_seqlens (torch.Tensor): - The cumulative sequence lengths of the input tensor. Default: `None`. - output_dtype (torch.dtype): - The dtype of the output tensor. Default: `torch.float`. - If `None`, the output dtype will be the same as the input dtype. - - Returns: - (I + A)^-1 with the same shape as A - """ - assert A.shape[-1] in [16, 32, 64] - output_dtype = A.dtype if output_dtype is None else output_dtype - - B, T, H, BT = A.shape - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - NT = len(chunk_indices) if cu_seqlens is not None else triton.cdiv(T, BT) - - Ai = torch.zeros_like(A, dtype=output_dtype) - if BT == 16: - merge_fn = solve_tril_16x16_kernel - elif BT == 32: - merge_fn = merge_16x16_to_32x32_inverse_kernel - elif BT == 64: - merge_fn = merge_16x16_to_64x64_inverse_kernel - - merge_fn[NT, B * H]( - A=A, - Ai=Ai, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - T=T, - H=H, - BT=BT, - USE_TMA=IS_TMA_SUPPORTED, - ) - return Ai diff --git a/src/xorl/ops/loss/__init__.py b/src/xorl/ops/loss/__init__.py index 18a8b620..cd9bd894 100644 --- a/src/xorl/ops/loss/__init__.py +++ b/src/xorl/ops/loss/__init__.py @@ -8,16 +8,23 @@ - drgrpo_loss_function: DR-GRPO loss with PPO clipping and KL penalty """ -from typing import Callable, Dict +from typing import Callable, Dict, Literal 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.opd_loss import OPDLossMetrics, opd_loss_function from xorl.ops.loss.policy_loss import policy_loss_function +from xorl.ops.loss.reducers import Reducer, SequencePartial, TokenPartial from xorl.ops.loss.vocab_parallel_cross_entropy import vocab_parallel_cross_entropy +# Cross-entropy computation mode shared by the local-trainer (TrainingArguments) +# and server-runner (ServerArguments) entry points so the Literal stays in sync. +CrossEntropyMode = Literal["eager", "compiled"] + + # --------------------------------------------------------------------------- # Loss function registry # --------------------------------------------------------------------------- @@ -27,6 +34,7 @@ "importance_sampling": importance_sampling_loss_function, "policy_loss": policy_loss_function, "drgrpo": drgrpo_loss_function, + "opd_loss": opd_loss_function, } @@ -43,13 +51,19 @@ def register_loss_function(name: str, fn: Callable) -> None: __all__ = [ + "CrossEntropyMode", "LossOutput", + "OPDLossMetrics", "LOSS_REGISTRY", + "Reducer", + "SequencePartial", + "TokenPartial", "get_loss_function", "register_loss_function", "causallm_loss_function", "drgrpo_loss_function", "importance_sampling_loss_function", + "opd_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 29bed248..75c212ff 100644 --- a/src/xorl/ops/loss/causallm_loss.py +++ b/src/xorl/ops/loss/causallm_loss.py @@ -3,8 +3,12 @@ import torch import torch.nn.functional as F -from xorl.ops.loss.compiled_cross_entropy import compiled_cross_entropy_function +from xorl.ops.loss.compiled_cross_entropy import ( + compiled_ce_and_lse_sq_function, + compiled_cross_entropy_function, +) from xorl.ops.loss.loss_output import LossOutput +from xorl.ops.loss.reducers import Reducer, TokenPartial from xorl.ops.loss.vocab_parallel_cross_entropy import vocab_parallel_cross_entropy @@ -19,6 +23,8 @@ def causallm_loss_function( tp_group=None, use_compile: bool = False, lm_head_fp32: bool = False, + loss_reducer: Reducer | None = None, + z_loss_coef: float = 0.0, ) -> "LossOutput": """ Compute causal language modeling loss. @@ -38,9 +44,25 @@ def causallm_loss_function( ce_mode: Cross-entropy mode - "compiled" (default) or "eager" num_chunks: Number of chunks for compiled mode (default: 8). tp_group: TP process group for vocab-parallel cross-entropy (default: None). + loss_reducer: Optional ``(values, mask) -> scalar``. When supplied, the + returned loss is a partial share under the reducer's denominator + (sum across micro-batches + all-reduce across ranks recovers the + globally-correct loss). When None, falls back to a local token mean. + Z-loss (when enabled) is reduced through the same reducer so the + two terms compose consistently. + z_loss_coef: If > 0, add the Z-loss auxiliary term used in OLMo / + PaLM-style training: + z_loss = coef * sum(logsumexp(logits)^2 * mask) / num_valid_tokens + where ``mask = labels != ignore_index``. Equivalent to OLMo's + ``cross_entropy_loss(..., reduction="sum")`` path divided by + ``batch_size_in_tokens``. Encourages log(Z) to stay near zero, + stabilizing training at large vocab / high LR. Not supported + in the TP path. Returns: LossOutput with loss, and optionally per_token_logprobs/per_token_loss. + When ``z_loss_coef > 0``, ``LossOutput.metrics`` contains + ``{"ce_loss": , "z_loss": }``. """ # Store original shape before flattening for per-token outputs original_shape = labels.shape @@ -50,8 +72,18 @@ def causallm_loss_function( hidden_states_flat = hidden_states.view(-1, hidden_states.size(-1)) valid_mask = labels_flat != ignore_index + if loss_reducer is None: + loss_reducer = TokenPartial(scale=valid_mask.sum().float()) + + mask_flat = valid_mask.float() + # Vocab-parallel cross-entropy for tensor parallelism if tp_group is not None: + if z_loss_coef > 0.0: + raise NotImplementedError( + "softmax_auxiliary_loss (Z-loss) is not yet supported with tensor parallelism. " + "Disable softmax_auxiliary_loss or run without TP." + ) # Extract local weight from DTensor if needed local_weight = weight.to_local() if hasattr(weight, "to_local") else weight @@ -64,7 +96,7 @@ def causallm_loss_function( use_compile=use_compile, ) - loss = per_token_ce.sum() / valid_mask.sum().clamp(min=1) + loss = loss_reducer(per_token_ce, mask_flat) if return_per_token: return LossOutput( loss=loss, @@ -73,40 +105,72 @@ def causallm_loss_function( ) return LossOutput(loss=loss) + z_loss_enabled = z_loss_coef > 0.0 + if return_per_token: - # Compute cross-entropy based on mode + # Compute cross-entropy based on mode (and Z-loss when enabled). + per_token_lse_sq = None 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 - ) + if z_loss_enabled: + per_token_ce, per_token_lse_sq = compiled_ce_and_lse_sq_function( + hidden_states_flat, weight, labels_flat, ignore_index, num_chunks, lm_head_fp32=lm_head_fp32 + ) + else: + 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() else: logits_flat = (hidden_states_flat @ weight.t()).float() per_token_ce = F.cross_entropy(logits_flat, labels_flat, reduction="none", ignore_index=ignore_index) + if z_loss_enabled: + lse = torch.logsumexp(logits_flat, dim=-1) + per_token_lse_sq = (lse * lse) * valid_mask.to(lse.dtype) - loss = per_token_ce.sum() / valid_mask.sum().clamp(min=1) + ce_loss = loss_reducer(per_token_ce, mask_flat) + if z_loss_enabled: + z_loss = loss_reducer(per_token_lse_sq, mask_flat) + loss = ce_loss + z_loss_coef * z_loss + metrics = {"ce_loss": ce_loss.detach(), "z_loss": z_loss.detach()} + else: + loss = ce_loss + metrics = None return LossOutput( loss=loss, per_token_logprobs=-per_token_ce.detach().view(original_shape), per_token_loss=per_token_ce.view(original_shape), + metrics=metrics, ) else: # Always use reduction="none" + manual mean to avoid NaN when all labels # are ignore_index (reduction="mean" returns NaN for 0 valid elements). # Keeping the autograd graph intact is critical for FSDP2: all ranks must # trigger reduce-scatter for every parameter, including lm_head weight. + per_token_lse_sq = None 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 - ) + if z_loss_enabled: + per_token_ce, per_token_lse_sq = compiled_ce_and_lse_sq_function( + hidden_states_flat, weight, labels_flat, ignore_index, num_chunks, lm_head_fp32=lm_head_fp32 + ) + else: + 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() else: logits_flat = (hidden_states_flat @ weight.t()).float() per_token_ce = F.cross_entropy(logits_flat, labels_flat, reduction="none", ignore_index=ignore_index) + if z_loss_enabled: + lse = torch.logsumexp(logits_flat, dim=-1) + per_token_lse_sq = (lse * lse) * valid_mask.to(lse.dtype) - loss = per_token_ce.sum() / valid_mask.sum().clamp(min=1) - return LossOutput(loss=loss) + ce_loss = loss_reducer(per_token_ce, mask_flat) + if z_loss_enabled: + z_loss = loss_reducer(per_token_lse_sq, mask_flat) + loss = ce_loss + z_loss_coef * z_loss + return LossOutput(loss=loss, metrics={"ce_loss": ce_loss.detach(), "z_loss": z_loss.detach()}) + return LossOutput(loss=ce_loss) diff --git a/src/xorl/ops/loss/compiled_cross_entropy.py b/src/xorl/ops/loss/compiled_cross_entropy.py index c4daf237..a62d8c58 100644 --- a/src/xorl/ops/loss/compiled_cross_entropy.py +++ b/src/xorl/ops/loss/compiled_cross_entropy.py @@ -14,6 +14,12 @@ # Cache for compiled cross-entropy functions _compiled_ce_cache: Dict[int, Callable] = {} +# Cache for compiled CE+LSE^2 (Z-loss) functions +_compiled_ce_and_lse_sq_cache: Dict[int, Callable] = {} + +# Cache for compiled OPD reverse-KL functions +_compiled_reverse_kl_cache: Dict[int, Callable] = {} + # Check if auto_chunker is available _AUTO_CHUNKER_AVAILABLE = None @@ -74,6 +80,101 @@ def compiled_cross_entropy_function( return compute_ce_fn(hidden_states, weight, labels, ignore_index) +def compiled_ce_and_lse_sq_function( + hidden_states: torch.Tensor, + weight: torch.Tensor, + labels: torch.Tensor, + ignore_index: int = -100, + num_chunks: int = 64, + lm_head_fp32: bool = False, +) -> tuple[torch.Tensor, torch.Tensor]: + """Compute per-token cross-entropy AND per-token logsumexp(logits)^2 in one fused pass. + + Used for the Z-loss auxiliary term. Without auto_chunker we'd have to + materialize the [batch*seq, vocab] logits tensor twice (once for CE, once + for LSE), so we co-compute them inside the same compiled region. + + Returns: + (per_token_ce, per_token_lse_sq) — both shape (batch * seq_len,). + per_token_lse_sq is zero at ignored-index positions. + """ + if lm_head_fp32: + hidden_states = hidden_states.float() + weight = weight.float() + fn = _get_compiled_ce_and_lse_sq_fn(num_chunks) + return fn(hidden_states, weight, labels, ignore_index) + + +def compiled_reverse_kl_function( + student_hidden_states: torch.Tensor, + student_weight: torch.Tensor, + teacher_hidden_states: torch.Tensor, + teacher_weight: torch.Tensor, + labels: torch.Tensor, + ignore_index: int = -100, + num_chunks: int = 64, + lm_head_fp32: bool = False, + teacher_lm_head_fp32: bool = True, +) -> torch.Tensor: + """Compute per-token KL(student || teacher) without materializing full logits twice. + + Args: + student_hidden_states: Flattened student hidden states, shape [tokens, student_hidden_dim]. + student_weight: Student LM head, shape [vocab_size, student_hidden_dim]. + teacher_hidden_states: Flattened teacher hidden states, shape [tokens, teacher_hidden_dim]. + teacher_weight: Teacher LM head, shape [vocab_size, teacher_hidden_dim]. + labels: Flattened labels, shape [tokens]. Only used to zero ignored positions. + ignore_index: Label value to mask out of the returned per-token KL. + num_chunks: Number of token chunks for torch.compile auto_chunker. 0 disables chunking. + lm_head_fp32: Cast student hidden/head tensors to fp32 before matmul. + teacher_lm_head_fp32: Cast teacher hidden/head tensors to fp32 before matmul. + + Returns: + Per-token reverse KL, shape [tokens], with zero at ignored-index positions. + """ + if lm_head_fp32: + student_hidden_states = student_hidden_states.float() + student_weight = student_weight.float() + if teacher_lm_head_fp32: + teacher_hidden_states = teacher_hidden_states.float() + teacher_weight = teacher_weight.float() + if student_hidden_states.device.type == "cpu": + return _compute_reverse_kl( + student_hidden_states, + student_weight, + teacher_hidden_states.detach(), + teacher_weight.detach(), + labels, + ignore_index, + ) + fn = _get_compiled_reverse_kl_fn(num_chunks) + return fn( + student_hidden_states, + student_weight, + teacher_hidden_states.detach(), + teacher_weight.detach(), + labels, + ignore_index, + ) + + +def _compute_reverse_kl( + student_hidden_states, + student_weight, + teacher_hidden_states, + teacher_weight, + labels, + ignore_index, +): + student_logits = (student_hidden_states @ student_weight.t()).float() + teacher_logits = (teacher_hidden_states @ teacher_weight.t()).float() + student_log_probs = F.log_softmax(student_logits, dim=-1) + teacher_log_probs = F.log_softmax(teacher_logits, dim=-1) + token_kl = (student_log_probs.exp() * (student_log_probs - teacher_log_probs)).sum(dim=-1) + valid = (labels != ignore_index).to(token_kl.dtype) + return token_kl * valid + + def _get_compiled_ce_fn(num_chunks: int, reduction: str = "none") -> Callable: """ Get or create a compiled cross-entropy function. @@ -107,3 +208,40 @@ def _compute_ce(hidden_states, weight, labels, ignore_index): else: _compiled_ce_cache[cache_key] = torch.compile(_compute_ce) return _compiled_ce_cache[cache_key] + + +def _get_compiled_ce_and_lse_sq_fn(num_chunks: int) -> Callable: + """Get or create a compiled CE+LSE^2 function (chunked along the token dim).""" + cache_key = num_chunks + if cache_key not in _compiled_ce_and_lse_sq_cache: + + def _compute_ce_and_lse_sq(hidden_states, weight, labels, ignore_index): + logits = (hidden_states @ weight.t()).float() + per_token_ce = F.cross_entropy(logits, labels, reduction="none", ignore_index=ignore_index) + lse = torch.logsumexp(logits, dim=-1) + valid = (labels != ignore_index).to(lse.dtype) + per_token_lse_sq = (lse * lse) * valid + return per_token_ce, per_token_lse_sq + + if num_chunks > 0 and _check_auto_chunker_available(): + _compiled_ce_and_lse_sq_cache[cache_key] = torch.compile( + _compute_ce_and_lse_sq, + options={"auto_chunker.enable": True, "auto_chunker.num_chunk": num_chunks}, + ) + else: + _compiled_ce_and_lse_sq_cache[cache_key] = torch.compile(_compute_ce_and_lse_sq) + return _compiled_ce_and_lse_sq_cache[cache_key] + + +def _get_compiled_reverse_kl_fn(num_chunks: int) -> Callable: + """Get or create a compiled reverse-KL function (chunked along the token dim).""" + cache_key = num_chunks + if cache_key not in _compiled_reverse_kl_cache: + if num_chunks > 0 and _check_auto_chunker_available(): + _compiled_reverse_kl_cache[cache_key] = torch.compile( + _compute_reverse_kl, + options={"auto_chunker.enable": True, "auto_chunker.num_chunk": num_chunks}, + ) + else: + _compiled_reverse_kl_cache[cache_key] = torch.compile(_compute_reverse_kl) + return _compiled_reverse_kl_cache[cache_key] diff --git a/src/xorl/ops/loss/grpo_loss.py b/src/xorl/ops/loss/grpo_loss.py index 8628a41f..14ce4912 100644 --- a/src/xorl/ops/loss/grpo_loss.py +++ b/src/xorl/ops/loss/grpo_loss.py @@ -5,38 +5,25 @@ https://arxiv.org/abs/2503.20783 """ -from typing import List, Literal, Optional, Tuple +from typing import List, Literal, 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 +from xorl.ops.loss.reducers import Reducer, TokenPartial -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, + metric_reducer: Reducer, ratio_type: RatioType = "token", ) -> Tuple[torch.Tensor, torch.Tensor, List[Tuple[str, torch.Tensor]]]: """Importance sampling ratio r = π_θ/π_old. @@ -60,8 +47,8 @@ def compute_ratio( with torch.no_grad(): metrics = [ - ("loss/ratio/mean", masked_mean(ratio, mask)), - ("loss/kl_policy/mean", masked_mean(-log_ratio, mask)), + ("loss/ratio/mean", metric_reducer(ratio, mask)), + ("loss/kl_policy/mean", metric_reducer(-log_ratio, mask)), ] return ratio, log_ratio, metrics @@ -71,6 +58,7 @@ def compute_kl( policy_logprobs: torch.Tensor, ref_logprobs: torch.Tensor, mask: torch.Tensor, + metric_reducer: Reducer, kl_type: KLType = "k3", ) -> Tuple[torch.Tensor, List[Tuple[str, torch.Tensor]]]: """KL divergence using Schulman's estimators (k1, k2, k3).""" @@ -88,7 +76,7 @@ def compute_kl( raise ValueError(f"Unknown kl_type: {kl_type}") with torch.no_grad(): - metrics = [("loss/kl_ref/mean", masked_mean(kl, mask))] + metrics = [("loss/kl_ref/mean", metric_reducer(kl, mask))] return kl, metrics @@ -97,6 +85,7 @@ def pg_ppo_clip( ratio: torch.Tensor, advantages: torch.Tensor, mask: torch.Tensor, + metric_reducer: Reducer, clip_low: float = 0.2, clip_high: float = 0.2, ) -> Tuple[torch.Tensor, List[Tuple[str, torch.Tensor]]]: @@ -114,62 +103,39 @@ def pg_ppo_clip( 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)), + ("loss/clip/clipped_ratio/mean", metric_reducer(clipped_ratio, mask)), + ("loss/clip/high_fraction", metric_reducer((clipped_high & pos_adv).float(), mask)), + ("loss/clip/low_fraction", metric_reducer((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, + ref_logprobs: torch.Tensor | None = 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, + tp_group: dist.ProcessGroup | None = None, lm_head_fp32: bool = False, - loss_scale: Optional[torch.Tensor] = None, + loss_reducer: Reducer | None = None, + metric_reducer: Reducer | None = 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) + Aggregated: ``loss_reducer(per_token_loss, mask)``. Defaults to + ``TokenPartial(scale=loss_mask.sum())`` — the local active-token mean. Args: hidden_states: (B, S, H) model hidden states. @@ -182,14 +148,16 @@ def drgrpo_loss_function( 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. + loss_reducer / metric_reducer: Both default to + ``TokenPartial(scale=loss_mask.sum())`` (legacy local active-token + mean; does not compose across mbs/ranks). Pass shared global-scale + reducers to make summed partial shares recover the global value. Returns: LossOutput with loss, per_token_logprobs, per_token_loss, and metrics. @@ -217,19 +185,23 @@ def drgrpo_loss_function( loss_mask = (labels != ignore_index).float() - ratio, _, ratio_m = compute_ratio(logprobs, old_logprobs, loss_mask, ratio_type) + if metric_reducer is None: + metric_reducer = TokenPartial(scale=loss_mask.sum()) + if loss_reducer is None: + loss_reducer = TokenPartial(scale=loss_mask.sum()) + + ratio, _, ratio_m = compute_ratio(logprobs, old_logprobs, loss_mask, metric_reducer, ratio_type) - pg_loss, clip_m = pg_ppo_clip(ratio, advantages, loss_mask, clip_low, clip_high) + pg_loss, clip_m = pg_ppo_clip(ratio, advantages, loss_mask, metric_reducer, 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) + kl, kl_m = compute_kl(logprobs, ref_logprobs, loss_mask, metric_reducer, kl_type) pg_loss = pg_loss + beta * kl - loss, agg_m = aggregate(pg_loss, loss_mask, agg_type, loss_scale) + loss = loss_reducer(pg_loss, loss_mask) - all_metrics = ratio_m + clip_m + kl_m + agg_m - metrics = {k: v.item() for k, v in all_metrics} + metrics = dict(ratio_m + clip_m + kl_m) return LossOutput( loss=loss, diff --git a/src/xorl/ops/loss/importance_sampling_loss.py b/src/xorl/ops/loss/importance_sampling_loss.py index beb5e274..b34f8b1e 100644 --- a/src/xorl/ops/loss/importance_sampling_loss.py +++ b/src/xorl/ops/loss/importance_sampling_loss.py @@ -1,12 +1,13 @@ from __future__ import annotations -from typing import Optional +from typing import Any, Dict, Optional 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 +from xorl.ops.loss.reducers import Reducer, TokenPartial def importance_sampling_loss_function( @@ -22,6 +23,8 @@ def importance_sampling_loss_function( tp_group: Optional[dist.ProcessGroup] = None, compute_kl_stats: bool = False, lm_head_fp32: bool = False, + loss_reducer: Optional[Reducer] = None, + metric_reducer: Optional[Reducer] = None, ) -> "LossOutput": """ Compute importance sampling loss for GRPO/RL training. @@ -50,6 +53,13 @@ def importance_sampling_loss_function( where log_ratio = new_logprobs - old_logprobs. Non-negative, unbiased, lower variance. - entropy_sample: -mean(old_logprobs) over valid tokens - valid_tokens: Count of valid tokens + loss_reducer: Reduces per-token loss to a scalar partial share. None => + ``TokenPartial(scale=valid_mask.sum())`` (legacy local token-mean; does + not compose across micro-batches/ranks). Pass a shared global-scale + reducer to make summed partial shares recover the global loss. + metric_reducer: Reduces per-token /mean metrics (ratio_mean, + kl_sample_train_k3, entropy_sample). ratio_min/ratio_max stay local + scalars and bypass it. Returns: LossOutput with loss, per_token_logprobs, per_token_loss, and metrics. @@ -65,7 +75,13 @@ def importance_sampling_loss_function( # Valid/action mask valid_mask = labels_flat != ignore_index - n_valid = valid_mask.sum().clamp(min=1).float() + valid_mask_f = valid_mask.float() + valid_count = valid_mask.sum() + + if loss_reducer is None: + loss_reducer = TokenPartial(scale=valid_count.float()) + if metric_reducer is None: + metric_reducer = TokenPartial(scale=valid_count.float()) # ---- Cross-entropy computation ---- per_token_ce = compute_per_token_ce( @@ -93,42 +109,34 @@ def importance_sampling_loss_function( per_token_pg = per_token_pg.masked_fill(~valid_mask, 0.0) # ---- Option B: value from true PG, grad from weighted CE surrogate ---- - true_pg = per_token_pg.sum() / n_valid + true_pg = loss_reducer(per_token_pg, valid_mask_f) w = (ratio.detach() * advantages_flat).masked_fill(~valid_mask, 0.0) - surrogate = (w * per_token_ce).sum() / n_valid + surrogate = loss_reducer(w * per_token_ce, valid_mask_f) loss = true_pg.detach() + surrogate - surrogate.detach() - # Compute metrics for logging (convert to Python floats for JSON serialization) - valid_ratio = ratio[valid_mask] if valid_mask.any() else ratio - metrics = { - "ratio_mean": valid_ratio.mean().detach().item(), - "ratio_min": valid_ratio.min().detach().item(), - "ratio_max": valid_ratio.max().detach().item(), + # ±inf identity on empty ranks lets cross-rank MIN/MAX-allreduce ignore empty contributors. + if valid_mask.any(): + ratio_min = ratio.masked_fill(~valid_mask, float("inf")).min() + ratio_max = ratio.masked_fill(~valid_mask, float("-inf")).max() + else: + ratio_min = ratio.new_tensor(float("inf")) + ratio_max = ratio.new_tensor(float("-inf")) + metrics: Dict[str, Any] = { + "ratio_mean": metric_reducer(ratio, valid_mask_f).detach(), + "ratio_min": ratio_min.detach(), + "ratio_max": ratio_max.detach(), } - # Optionally compute KL statistics if compute_kl_stats: with torch.no_grad(): - _n_valid_kl = valid_mask.sum().item() # TRUE count, no clamp - if valid_mask.any(): - valid_old = old_logprobs_flat[valid_mask] - valid_new = new_logprobs_flat[valid_mask] - log_ratio = valid_new - valid_old - - # K3 estimator (Schulman): exp(log_ratio) - log_ratio - 1 - # Non-negative, unbiased, lower variance than K1/K2 - k3 = (torch.exp(log_ratio) - log_ratio - 1.0).mean().item() - metrics["kl_sample_train_k3"] = k3 - metrics["entropy_sample"] = -valid_old.mean().item() - metrics["valid_tokens"] = _n_valid_kl - metrics["_n_valid_kl"] = _n_valid_kl - else: - metrics["kl_sample_train_k3"] = 0.0 - metrics["entropy_sample"] = 0.0 - metrics["valid_tokens"] = 0 - metrics["_n_valid_kl"] = 0 + log_ratio_full = (new_logprobs_flat - old_logprobs_flat).masked_fill(~valid_mask, 0.0) + ratio_full = torch.exp(log_ratio_full) + per_token_k3 = ratio_full - log_ratio_full - 1.0 + metrics["kl_sample_train_k3"] = metric_reducer(per_token_k3, valid_mask_f) + metrics["entropy_sample"] = metric_reducer(-old_logprobs_flat, valid_mask_f) + metrics["valid_tokens"] = valid_count.item() # Reshape per-token outputs per_token_logprobs = new_logprobs_flat.view(original_shape) @@ -139,4 +147,5 @@ def importance_sampling_loss_function( per_token_logprobs=per_token_logprobs, per_token_loss=per_token_loss, metrics=metrics, + metric_ops={"ratio_min": "min", "ratio_max": "max"}, ) diff --git a/src/xorl/ops/loss/loss_output.py b/src/xorl/ops/loss/loss_output.py index c1c8e588..41a79eee 100644 --- a/src/xorl/ops/loss/loss_output.py +++ b/src/xorl/ops/loss/loss_output.py @@ -6,9 +6,16 @@ @dataclass class LossOutput: - """Standardized return type for all loss functions.""" + """Standardized return type for all loss functions. + + ``metric_ops`` tags ``metrics`` keys whose cross-mb / cross-rank composition + isn't the default mean (``"min"``/``"max"``). The sidecar (rather than a + tagged-value type in ``metrics``) keeps the metrics dict directly + JSON-serializable for untagged consumers. + """ loss: torch.Tensor per_token_logprobs: Optional[torch.Tensor] = None per_token_loss: Optional[torch.Tensor] = None metrics: Optional[Dict[str, Any]] = None + metric_ops: Optional[Dict[str, str]] = None diff --git a/src/xorl/ops/loss/opd_loss.py b/src/xorl/ops/loss/opd_loss.py new file mode 100644 index 00000000..b916483e --- /dev/null +++ b/src/xorl/ops/loss/opd_loss.py @@ -0,0 +1,214 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Optional + +import torch + +from xorl.ops.loss.compiled_cross_entropy import compiled_reverse_kl_function +from xorl.ops.loss.loss_output import LossOutput +from xorl.ops.loss.opd_streaming_kl import streaming_reverse_kl_function +from xorl.ops.loss.reducers import Reducer, TokenPartial + + +@dataclass(frozen=True) +class OPDLossMetrics: + valid_tokens: int + opd_kl: float = 0.0 + opd_weighted_kl: float = 0.0 + opd_teacher_weight_mean: float = 0.0 + opd_num_teachers: Optional[int] = None + + def to_dict(self) -> dict[str, int | float]: + metrics: dict[str, int | float] = { + "valid_tokens": self.valid_tokens, + "opd_kl": self.opd_kl, + "opd_weighted_kl": self.opd_weighted_kl, + "opd_teacher_weight_mean": self.opd_teacher_weight_mean, + } + if self.opd_num_teachers is not None: + metrics["opd_num_teachers"] = self.opd_num_teachers + return metrics + + +def _as_flat_optional_weights( + teacher_weights: Optional[torch.Tensor], + valid_mask: torch.Tensor, + dtype: torch.dtype, +) -> torch.Tensor: + if teacher_weights is None: + return torch.ones(valid_mask.sum(), dtype=dtype, device=valid_mask.device) + weights_flat = teacher_weights.reshape(-1).to(device=valid_mask.device, dtype=dtype) + return weights_flat[valid_mask] + + +def _zero_loss_with_graph(hidden_states: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + """Build a 0-valued loss that still flows gradients through hidden_states + weight. + + Always returns fp32 so the dtype matches the normal-return path + (`total_weighted_kl / denom`, fp32). A dtype mismatch between the early-return + branch on no-valid-token ranks and the fp32 normal branch corrupts NCCL + all_reduce in the trainer's loss-reporting path. + """ + return hidden_states.float().sum() * 0.0 + weight.float().sum() * 0.0 + + +def _denominator_tensor( + denominator: torch.Tensor | int | float | None, + *, + fallback: torch.Tensor, + device: torch.device | str, +) -> torch.Tensor: + if denominator is None: + return fallback.to(device=device, dtype=torch.float32) + if torch.is_tensor(denominator): + return denominator.to(device=device, dtype=torch.float32) + return torch.tensor(float(denominator), device=device, dtype=torch.float32) + + +def opd_loss_function( + hidden_states: torch.Tensor, + weight: torch.Tensor, + labels: torch.Tensor, + teacher_hidden_states: torch.Tensor, + teacher_lm_head_weight: torch.Tensor, + teacher_weights: Optional[torch.Tensor] = None, + ignore_index: int = -100, + num_chunks: int = 8, + lm_head_fp32: bool = False, + teacher_lm_head_fp32: bool = True, + kl_backend: str = "torch_compile", + vocab_chunk_size: int = 32768, + return_per_token: bool = False, + normalization_denominator: Optional[torch.Tensor | int | float] = None, + loss_reducer: Optional[Reducer] = None, + metric_reducer: Optional[Reducer] = None, +) -> LossOutput: + """Compute full-vocabulary reverse KL for on-policy distillation. + + The objective is KL(student || teacher) at each valid token position: + sum_v p_student(v) * (log p_student(v) - log p_teacher(v)). + + Expected shapes: + hidden_states: [batch, seq, student_hidden_dim] + weight: [vocab_size, student_hidden_dim] + labels: [batch, seq], with ignore_index masking tokens out of the loss + teacher_hidden_states: [batch, seq, teacher_hidden_dim] + teacher_lm_head_weight: [vocab_size, teacher_hidden_dim] + teacher_weights: optional [batch, seq] per-token multipliers applied + after KL computation and before the final normalization. These are + useful for mixing teachers or down-weighting lower-confidence teacher + outputs without changing the valid-token denominator. + + Teacher tensors are detached by construction. Only the student hidden states + and student LM head receive gradients. + + ``kl_backend`` selects the full-vocabulary KL implementation. ``torch_compile`` + preserves the existing auto-chunked path. ``streaming`` and ``tilelang`` use + the OPD streaming path that saves only per-token normalization statistics and + recomputes vocab chunks in backward; ``tilelang`` is the stable selector for + the future native TileLang kernel. + """ + if hidden_states.shape[:-1] != labels.shape: + raise ValueError(f"hidden_states shape {hidden_states.shape} is incompatible with labels {labels.shape}") + if teacher_hidden_states.shape[:-1] != labels.shape: + raise ValueError( + f"teacher_hidden_states shape {teacher_hidden_states.shape} is incompatible with labels {labels.shape}" + ) + if weight.shape[0] != teacher_lm_head_weight.shape[0]: + raise ValueError( + f"student vocab size ({weight.shape[0]}) must match teacher vocab size ({teacher_lm_head_weight.shape[0]})" + ) + if hidden_states.shape[-1] != weight.shape[-1]: + raise ValueError( + f"student hidden size ({hidden_states.shape[-1]}) must match student head width ({weight.shape[-1]})" + ) + if teacher_hidden_states.shape[-1] != teacher_lm_head_weight.shape[-1]: + raise ValueError( + "teacher hidden size " + f"({teacher_hidden_states.shape[-1]}) must match teacher head width ({teacher_lm_head_weight.shape[-1]})" + ) + + original_shape = labels.shape + labels_flat = labels.reshape(-1) + valid_mask = labels_flat != ignore_index + valid_count = valid_mask.sum() + + if valid_count.item() == 0: + loss = _zero_loss_with_graph(hidden_states, weight) + per_token_loss = ( + torch.zeros(original_shape, dtype=torch.float32, device=labels.device) if return_per_token else None + ) + return LossOutput(loss=loss, per_token_loss=per_token_loss, metrics=OPDLossMetrics(valid_tokens=0).to_dict()) + + student_hidden_flat = hidden_states.reshape(-1, hidden_states.size(-1))[valid_mask] + teacher_hidden_flat = teacher_hidden_states.reshape(-1, teacher_hidden_states.size(-1))[valid_mask].detach() + labels_valid = labels_flat[valid_mask] + token_weights = _as_flat_optional_weights(teacher_weights, valid_mask, torch.float32) + + default_scale = _denominator_tensor( + normalization_denominator, + fallback=valid_count, + device=hidden_states.device, + ) + if loss_reducer is None: + loss_reducer = TokenPartial(scale=default_scale) + if metric_reducer is None: + metric_reducer = TokenPartial(scale=default_scale) + + backend = kl_backend.lower() + if backend in {"torch_compile", "compile", "auto_chunker"}: + if not torch.is_tensor(teacher_lm_head_weight): + raise ValueError("torch_compile OPD KL backend requires a materialized teacher LM head tensor") + token_kl = compiled_reverse_kl_function( + student_hidden_states=student_hidden_flat, + student_weight=weight, + teacher_hidden_states=teacher_hidden_flat, + teacher_weight=teacher_lm_head_weight, + labels=labels_valid, + ignore_index=ignore_index, + num_chunks=num_chunks, + lm_head_fp32=lm_head_fp32, + teacher_lm_head_fp32=teacher_lm_head_fp32, + ) + elif backend in {"streaming", "tilelang"}: + if lm_head_fp32: + student_hidden_flat = student_hidden_flat.float() + weight = weight.float() + if teacher_lm_head_fp32: + teacher_hidden_flat = teacher_hidden_flat.float() + if torch.is_tensor(teacher_lm_head_weight): + teacher_lm_head_weight = teacher_lm_head_weight.float() + token_kl = streaming_reverse_kl_function( + student_hidden_states=student_hidden_flat, + student_weight=weight, + teacher_hidden_states=teacher_hidden_flat, + teacher_weight=teacher_lm_head_weight, + labels=labels_valid, + ignore_index=ignore_index, + vocab_chunk_size=vocab_chunk_size, + ) + else: + raise ValueError( + f"Unsupported OPD KL backend '{kl_backend}'. Expected 'torch_compile', 'streaming', or 'tilelang'." + ) + weighted_token_kl = token_kl * token_weights.to(token_kl.device) + valid_ones = torch.ones_like(weighted_token_kl, dtype=torch.float32) + + loss = loss_reducer(weighted_token_kl, valid_ones) + + per_token_loss = None + if return_per_token: + per_token_flat = torch.zeros(labels_flat.shape, dtype=torch.float32, device=labels.device) + per_token_flat[valid_mask] = weighted_token_kl.detach().to(per_token_flat.device) + per_token_loss = per_token_flat.view(original_shape) + + valid_count_float = max(float(valid_count.item()), 1.0) + metrics = OPDLossMetrics( + valid_tokens=int(valid_count.item()), + opd_kl=token_kl.detach().sum().item() / valid_count_float, + opd_weighted_kl=metric_reducer(weighted_token_kl.detach(), valid_ones).item(), + opd_teacher_weight_mean=token_weights.mean().item(), + ).to_dict() + + return LossOutput(loss=loss, per_token_loss=per_token_loss, metrics=metrics) diff --git a/src/xorl/ops/loss/opd_streaming_kl.py b/src/xorl/ops/loss/opd_streaming_kl.py new file mode 100644 index 00000000..ae6a66e7 --- /dev/null +++ b/src/xorl/ops/loss/opd_streaming_kl.py @@ -0,0 +1,188 @@ +from __future__ import annotations + +import math + +import torch + + +def _iter_weight_chunks(teacher_weight, vocab_size: int, chunk_rows: int): + if hasattr(teacher_weight, "iter_device_chunks"): + yield from teacher_weight.iter_device_chunks(chunk_rows) + return + for start, end in _iter_ranges(vocab_size, chunk_rows): + yield start, end, teacher_weight[start:end] + + +def _chunk_size(vocab_size: int, requested: int) -> int: + if requested <= 0 or requested >= vocab_size: + return vocab_size + return requested + + +def _iter_ranges(vocab_size: int, requested_chunk_size: int): + chunk_size = _chunk_size(vocab_size, requested_chunk_size) + for start in range(0, vocab_size, chunk_size): + yield start, min(start + chunk_size, vocab_size) + + +def _update_online_logsumexp( + running_max: torch.Tensor, + running_sumexp: torch.Tensor, + logits: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + chunk_max = logits.max(dim=-1, keepdim=True).values + new_max = torch.maximum(running_max, chunk_max) + prev_scale = torch.where( + torch.isfinite(running_max), + (running_max - new_max).exp(), + torch.zeros_like(running_sumexp), + ) + chunk_sumexp = (logits - new_max).exp().sum(dim=-1, keepdim=True) + return new_max, running_sumexp * prev_scale + chunk_sumexp + + +class _StreamingReverseKL(torch.autograd.Function): + """Exact KL(student || teacher) over vocab chunks. + + This is the TileLang-facing OPD path: it exposes the same execution shape a + native kernel will use (stream vocab blocks, save only per-token statistics, + recompute logits in backward) while keeping a pure PyTorch implementation as + the portable fallback. + """ + + @staticmethod + def forward( + ctx, + student_hidden_states: torch.Tensor, + student_weight: torch.Tensor, + teacher_hidden_states: torch.Tensor, + labels: torch.Tensor, + teacher_weight, + ignore_index: int, + vocab_chunk_size: int, + ) -> torch.Tensor: + teacher_shape = tuple(int(x) for x in teacher_weight.shape) + if student_weight.shape[0] != teacher_shape[0]: + raise ValueError( + f"student vocab size ({student_weight.shape[0]}) must match teacher vocab size ({teacher_shape[0]})" + ) + + vocab_size = int(student_weight.shape[0]) + token_count = int(student_hidden_states.shape[0]) + valid = labels != ignore_index + neg_inf = -float("inf") + + s_max = torch.full((token_count, 1), neg_inf, device=student_hidden_states.device, dtype=torch.float32) + t_max = torch.full((token_count, 1), neg_inf, device=student_hidden_states.device, dtype=torch.float32) + s_sumexp = torch.zeros((token_count, 1), device=student_hidden_states.device, dtype=torch.float32) + t_sumexp = torch.zeros((token_count, 1), device=student_hidden_states.device, dtype=torch.float32) + + for start, end, t_weight in _iter_weight_chunks(teacher_weight, vocab_size, vocab_chunk_size): + s_logits = (student_hidden_states @ student_weight[start:end].t()).float() + t_logits = (teacher_hidden_states @ t_weight.t()).float() + s_max, s_sumexp = _update_online_logsumexp(s_max, s_sumexp, s_logits) + t_max, t_sumexp = _update_online_logsumexp(t_max, t_sumexp, t_logits) + + s_logz = s_sumexp.log() + s_max + t_logz = t_sumexp.log() + t_max + kl = torch.zeros(token_count, device=student_hidden_states.device, dtype=torch.float32) + for start, end, t_weight in _iter_weight_chunks(teacher_weight, vocab_size, vocab_chunk_size): + s_logits = (student_hidden_states @ student_weight[start:end].t()).float() + t_logits = (teacher_hidden_states @ t_weight.t()).float() + s_log_probs = s_logits - s_logz + t_log_probs = t_logits - t_logz + s_probs = s_log_probs.exp() + kl = kl + (s_probs * (s_log_probs - t_log_probs)).sum(dim=-1) + + kl = kl * valid.to(kl.dtype) + ctx.save_for_backward( + student_hidden_states, + student_weight, + teacher_hidden_states, + labels, + s_logz, + t_logz, + kl, + ) + ctx.teacher_weight = teacher_weight + ctx.ignore_index = ignore_index + ctx.vocab_chunk_size = vocab_chunk_size + return kl + + @staticmethod + def backward(ctx, grad_output: torch.Tensor): + ( + student_hidden_states, + student_weight, + teacher_hidden_states, + labels, + s_logz, + t_logz, + kl, + ) = ctx.saved_tensors + teacher_weight = ctx.teacher_weight + + valid = (labels != ctx.ignore_index).to(dtype=torch.float32, device=grad_output.device) + scale = grad_output.to(dtype=torch.float32) * valid + vocab_size = int(student_weight.shape[0]) + + grad_hidden = None + if ctx.needs_input_grad[0]: + grad_hidden = torch.zeros_like(student_hidden_states, dtype=torch.float32) + + grad_weight = None + if ctx.needs_input_grad[1]: + grad_weight = torch.zeros_like(student_weight, dtype=torch.float32) + + for start, end, t_weight in _iter_weight_chunks(teacher_weight, vocab_size, ctx.vocab_chunk_size): + s_weight = student_weight[start:end] + s_logits = (student_hidden_states @ s_weight.t()).float() + t_logits = (teacher_hidden_states @ t_weight.t()).float() + s_log_probs = s_logits - s_logz + t_log_probs = t_logits - t_logz + s_probs = s_log_probs.exp() + + # d KL(p_s || p_t) / d student_logits_i = + # p_s_i * (log p_s_i - log p_t_i - KL) + grad_logits = s_probs * (s_log_probs - t_log_probs - kl.unsqueeze(1)) + grad_logits = grad_logits * scale.unsqueeze(1) + + if grad_hidden is not None: + grad_hidden = grad_hidden + grad_logits @ s_weight.float() + if grad_weight is not None: + grad_weight[start:end] = grad_logits.t() @ student_hidden_states.float() + + if grad_hidden is not None: + grad_hidden = grad_hidden.to(student_hidden_states.dtype) + if grad_weight is not None: + grad_weight = grad_weight.to(student_weight.dtype) + + if hasattr(teacher_weight, "clear_device_cache"): + teacher_weight.clear_device_cache() + + return grad_hidden, grad_weight, None, None, None, None, None + + +def streaming_reverse_kl_function( + student_hidden_states: torch.Tensor, + student_weight: torch.Tensor, + teacher_hidden_states: torch.Tensor, + teacher_weight: torch.Tensor, + labels: torch.Tensor, + ignore_index: int = -100, + vocab_chunk_size: int = 32768, +) -> torch.Tensor: + """Compute per-token reverse KL without materializing full-vocab logits.""" + if vocab_chunk_size <= 0: + vocab_chunk_size = int(student_weight.shape[0]) + if not math.isfinite(float(vocab_chunk_size)): + raise ValueError(f"Invalid vocab_chunk_size={vocab_chunk_size}") + return _StreamingReverseKL.apply( + student_hidden_states, + student_weight, + teacher_hidden_states.detach(), + labels, + teacher_weight.detach() if torch.is_tensor(teacher_weight) else teacher_weight, + int(ignore_index), + int(vocab_chunk_size), + ) diff --git a/src/xorl/ops/loss/policy_loss.py b/src/xorl/ops/loss/policy_loss.py index 5616b5e9..d5ab3d09 100644 --- a/src/xorl/ops/loss/policy_loss.py +++ b/src/xorl/ops/loss/policy_loss.py @@ -10,13 +10,14 @@ from __future__ import annotations import logging -from typing import Dict, Optional, Tuple +from typing import Any, Dict, 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 +from xorl.ops.loss.reducers import Reducer, TokenPartial logger = logging.getLogger(__name__) @@ -71,6 +72,7 @@ def apply_tis_correction( train_log_probs: torch.Tensor, rollout_log_probs: torch.Tensor, valid_mask: torch.Tensor, + metric_reducer: Reducer, tis_clip_low: float = 0.1, tis_clip_high: float = 2.0, ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]: @@ -85,6 +87,8 @@ def apply_tis_correction( train_log_probs: Log probabilities from current training step rollout_log_probs: Log probabilities from rollout/inference valid_mask: Mask for valid tokens + metric_reducer: Reducer applied to per-token mean metrics (tis_mean, + tis_clipfrac). min/max are local reductions and bypass it. tis_clip_low: Lower bound for TIS clipping (default: 0.1) tis_clip_high: Upper bound for TIS clipping (default: 2.0) @@ -100,12 +104,20 @@ def apply_tis_correction( # Apply TIS correction to loss corrected_loss = pg_loss * tis_clipped - # Compute metrics + valid_mask_f = valid_mask.float() + tis_clipfrac_per_token = (tis_clipped != tis).float() + # ±inf identity on empty ranks lets cross-rank MIN/MAX-allreduce ignore empty contributors. + if valid_mask.any(): + tis_min = tis.masked_fill(~valid_mask, float("inf")).min() + tis_max = tis.masked_fill(~valid_mask, float("-inf")).max() + else: + tis_min = tis.new_tensor(float("inf")) + tis_max = tis.new_tensor(float("-inf")) tis_metrics = { - "tis_mean": tis[valid_mask].mean() if valid_mask.any() else torch.tensor(1.0), - "tis_min": tis[valid_mask].min() if valid_mask.any() else torch.tensor(1.0), - "tis_max": tis[valid_mask].max() if valid_mask.any() else torch.tensor(1.0), - "tis_clipfrac": (tis_clipped != tis)[valid_mask].float().mean() if valid_mask.any() else torch.tensor(0.0), + "tis_mean": metric_reducer(tis, valid_mask_f), + "tis_min": tis_min, + "tis_max": tis_max, + "tis_clipfrac": metric_reducer(tis_clipfrac_per_token, valid_mask_f), } return corrected_loss, tis_metrics @@ -132,6 +144,8 @@ def policy_loss_function( tp_group: Optional[dist.ProcessGroup] = None, lm_head_fp32: bool = False, icepop_beta: Optional[float] = None, + loss_reducer: Optional[Reducer] = None, + metric_reducer: Optional[Reducer] = None, ) -> "LossOutput": """ Policy loss with PPO clipping, optional IcePop masking, and optional TIS correction. @@ -166,6 +180,13 @@ def policy_loss_function( compute_kl_stats: If True, compute and return full KL statistics in metrics dict (kl_sample_train_k3, entropy_sample, ratio stats). If False (default), only return valid_tokens and pg_clipfrac. + loss_reducer: Reduces per-token loss to a scalar partial share. None => + ``TokenPartial(scale=valid_mask.sum())`` (legacy local token-mean; does + not compose across micro-batches/ranks). Pass a shared global-scale + reducer to make summed partial shares recover the global loss. + metric_reducer: Reduces per-token /mean metrics (pg_clipfrac, icepop_maskfrac, + tis_mean, tis_clipfrac, kl_sample_train_k3, entropy_sample, ratio_mean) + the same way. ratio_min/ratio_max/tis_min/tis_max stay local scalars. Returns: LossOutput with loss, per_token_logprobs (new logprobs), and metrics. @@ -182,7 +203,13 @@ def policy_loss_function( # Create mask for valid tokens (use labels != ignore_index) valid_mask = labels_flat != ignore_index - n_valid = valid_mask.sum().clamp(min=1) + valid_mask_f = valid_mask.float() + valid_count = valid_mask.sum() + + if loss_reducer is None: + loss_reducer = TokenPartial(scale=valid_count.float()) + if metric_reducer is None: + metric_reducer = TokenPartial(scale=valid_count.float()) # Compute cross-entropy (supports vocab-parallel TP via tp_group) per_token_ce = compute_per_token_ce( @@ -205,33 +232,27 @@ def policy_loss_function( ppo_kl = ppo_kl.masked_fill(~valid_mask, 0.0) advantages_masked = advantages_flat.masked_fill(~valid_mask, 0.0) - # Compute KL stats BEFORE compute_ppo_loss to avoid torch.compile interference + # Computed BEFORE compute_ppo_loss to avoid torch.compile interference. _kl_stats = None if compute_kl_stats: with torch.no_grad(): - _n_valid_kl = valid_mask.sum().item() # TRUE count, no clamp + _log_ratio_full = (new_logprobs_flat - old_logprobs_flat).masked_fill(~valid_mask, 0.0) + _ratio_full = torch.exp(_log_ratio_full) + _per_token_k3 = _ratio_full - _log_ratio_full - 1.0 + # ±inf identity on empty ranks lets cross-rank MIN/MAX-allreduce ignore empty contributors. if valid_mask.any(): - _valid_old = old_logprobs_flat[valid_mask] - _valid_new = new_logprobs_flat[valid_mask] - _log_ratio = _valid_new - _valid_old - _ratio_valid = torch.exp(_log_ratio) - _kl_stats = { - "kl_sample_train_k3": (_ratio_valid - _log_ratio - 1.0).mean().item(), - "entropy_sample": -_valid_old.mean().item(), - "ratio_mean": _ratio_valid.mean().item(), - "ratio_min": _ratio_valid.min().item(), - "ratio_max": _ratio_valid.max().item(), - "_n_valid_kl": _n_valid_kl, - } + _ratio_min = _ratio_full.masked_fill(~valid_mask, float("inf")).min() + _ratio_max = _ratio_full.masked_fill(~valid_mask, float("-inf")).max() else: - _kl_stats = { - "kl_sample_train_k3": 0.0, - "entropy_sample": 0.0, - "ratio_mean": 1.0, - "ratio_min": 1.0, - "ratio_max": 1.0, - "_n_valid_kl": 0, - } + _ratio_min = _ratio_full.new_tensor(float("inf")) + _ratio_max = _ratio_full.new_tensor(float("-inf")) + _kl_stats = { + "kl_sample_train_k3": metric_reducer(_per_token_k3, valid_mask_f), + "entropy_sample": metric_reducer(-old_logprobs_flat, valid_mask_f), + "ratio_mean": metric_reducer(_ratio_full, valid_mask_f), + "ratio_min": _ratio_min, + "ratio_max": _ratio_max, + } # Compute PPO-style clipped loss (returns per-token losses, clip mask, and ratio) pg_losses, is_clipped, ratio = compute_ppo_loss( @@ -263,13 +284,13 @@ def policy_loss_function( train_log_probs=new_logprobs_flat, rollout_log_probs=rollout_logprobs_flat, valid_mask=valid_mask, + metric_reducer=metric_reducer, tis_clip_low=tis_clip_low, tis_clip_high=tis_clip_high, ) - # Surrogate loss for gradient flow through CE - # True loss value (for logging): pg_losses averaged over valid tokens - true_loss = pg_losses.masked_fill(~valid_mask, 0.0).sum() / n_valid + # True loss value (for logging): partial share under loss_reducer. + true_loss = loss_reducer(pg_losses, valid_mask_f) # Gradient-active mask: tokens that are not clipped, not IcePop-masked, and valid gradient_active = ~is_clipped & valid_mask @@ -278,7 +299,7 @@ def policy_loss_function( # Surrogate: gradient weight = ratio * advantages, zeroed for inactive tokens gradient_weight = (ratio.detach() * advantages_flat).masked_fill(~gradient_active, 0.0) - surrogate = (gradient_weight * per_token_ce).sum() / n_valid + surrogate = loss_reducer(gradient_weight * per_token_ce, valid_mask_f) # Combine: forward value from true_loss, gradient from surrogate loss_with_grad = true_loss.detach() + surrogate - surrogate.detach() @@ -286,26 +307,31 @@ def policy_loss_function( # Return training logprobs reshaped new_logprobs = new_logprobs_flat.view(original_shape) - # Compute metrics with torch.no_grad(): - metrics = { - "valid_tokens": valid_mask.sum().item(), - "pg_clipfrac": is_clipped[valid_mask].float().mean().item() if valid_mask.any() else 0.0, + metrics: Dict[str, Any] = { + "valid_tokens": valid_count.item(), + "pg_clipfrac": metric_reducer(is_clipped.float(), valid_mask_f), } if icepop_mask is not None: - metrics["icepop_maskfrac"] = (~icepop_mask)[valid_mask].float().mean().item() if valid_mask.any() else 0.0 + metrics["icepop_maskfrac"] = metric_reducer((~icepop_mask).float(), valid_mask_f) - # Use pre-computed KL statistics (computed before torch.compile'd compute_ppo_loss) if _kl_stats is not None: metrics.update(_kl_stats) - # Add TIS metrics if available - for k, v in tis_metrics.items(): - metrics[k] = v.item() if torch.is_tensor(v) else v + metrics.update(tis_metrics) + + metric_ops: Dict[str, str] = {} + if _kl_stats is not None: + metric_ops["ratio_min"] = "min" + metric_ops["ratio_max"] = "max" + if tis_metrics: + metric_ops["tis_min"] = "min" + metric_ops["tis_max"] = "max" return LossOutput( loss=loss_with_grad, per_token_logprobs=new_logprobs, metrics=metrics, + metric_ops=metric_ops or None, ) diff --git a/src/xorl/ops/loss/reducers.py b/src/xorl/ops/loss/reducers.py new file mode 100644 index 00000000..48b414c8 --- /dev/null +++ b/src/xorl/ops/loss/reducers.py @@ -0,0 +1,71 @@ +"""Reducer protocol and canonical denominator policies for loss aggregation. + +A ``Reducer`` collapses a ``(B, S)`` tensor to a scalar over a +caller-supplied denominator policy. Partial shares sum across micro-batches +and ``all_reduce(SUM)`` across ranks to the globally-correct value. + +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Protocol, runtime_checkable + +import torch + + +@runtime_checkable +class Reducer(Protocol): + """``(values, mask) -> scalar`` partial share over a pre-computed + denominator. Partial shares sum across micro-batches and ``all_reduce(SUM)`` + across ranks. + """ + + def __call__(self, values: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: ... + + +@dataclass(frozen=True) +class TokenPartial: + """Flat masked sum divided by a caller-supplied ``scale``.""" + + scale: torch.Tensor + + def __call__(self, values: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: + return (values * mask).sum() / self.scale.clamp(min=1.0) + + +@dataclass(frozen=True) +class SequencePartial: + """Sum of per-segment token-means, divided by a caller-supplied ``scale``. + + Segment boundaries are flat across ``(values * mask).reshape(-1)``: + + - ``cu_seqlens_local: (N+1,)`` — shard-local segment extents. Under CP each + rank's slice sums to its segment's local contribution. + - ``seq_lengths_global: (N,)`` — pre-CP-shard token count per segment, used + as the per-segment denominator so partial shares from each CP rank sum + to the correct per-segment mean. + """ + + scale: torch.Tensor + cu_seqlens_local: torch.Tensor + seq_lengths_global: torch.Tensor + + def __call__(self, values: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: + flat = (values * mask).reshape(-1) + seg_lengths_local = self.cu_seqlens_local.diff() + n_segments = seg_lengths_local.numel() + seg_ids = torch.repeat_interleave( + torch.arange(n_segments, device=flat.device), + seg_lengths_local, + ) + seg_sums = torch.zeros(n_segments, dtype=flat.dtype, device=flat.device).index_add(0, seg_ids, flat) + seg_means = seg_sums / self.seq_lengths_global.clamp(min=1.0) + return seg_means.sum() / self.scale.clamp(min=1.0) + + +__all__ = [ + "Reducer", + "SequencePartial", + "TokenPartial", +] diff --git a/src/xorl/ops/moe/activations.py b/src/xorl/ops/moe/activations.py new file mode 100644 index 00000000..7e74e363 --- /dev/null +++ b/src/xorl/ops/moe/activations.py @@ -0,0 +1,107 @@ +"""MoE activation registry. + +A single source of truth for the activation functions used by all MoE +backends (native / eager / triton / quack). Each activation is a binary op +``(gate_out, up_out) -> h`` whose output is fed to the down projection. + +New activations are added by: + +1. Writing the implementation as a function of ``(gate_out, up_out)``. +2. Registering it in ``MOE_ACTIVATIONS`` under the canonical name. +3. Extending ``normalize_hidden_act`` to recognize any HF aliases. +4. Adding the name to the ``SUPPORTED_HIDDEN_ACTS`` set of each backend + whose kernel actually implements it (backends validate at entry). +""" + +from __future__ import annotations + +from typing import Callable, Dict + +import torch +import torch.nn.functional as F + + +# --------------------------------------------------------------------------- +# Activation constants +# --------------------------------------------------------------------------- + +# GPT-OSS clamped SwiGLU. If a future model needs different values, register +# a new ``hidden_act`` kind rather than making these runtime-configurable. +CLAMPED_SWIGLU_ALPHA: float = 1.702 +CLAMPED_SWIGLU_LIMIT: float = 7.0 + + +# --------------------------------------------------------------------------- +# Activation implementations +# --------------------------------------------------------------------------- + + +def silu_swiglu(gate_out: torch.Tensor, up_out: torch.Tensor) -> torch.Tensor: + """Standard SwiGLU: ``silu(gate) * up``.""" + return F.silu(gate_out) * up_out + + +def gelu_tanh_glu(gate_out: torch.Tensor, up_out: torch.Tensor) -> torch.Tensor: + """GeGLU (tanh approx): ``gelu_tanh(gate) * up``.""" + return F.gelu(gate_out, approximate="tanh") * up_out + + +def clamped_swiglu(gate_out: torch.Tensor, up_out: torch.Tensor) -> torch.Tensor: + """GPT-OSS clamped SwiGLU. + + Clamp both branches, then ``silu(alpha * gate) * (up + 1)``. + """ + gate_out = gate_out.clamp(max=CLAMPED_SWIGLU_LIMIT) + up_out = up_out.clamp(min=-CLAMPED_SWIGLU_LIMIT, max=CLAMPED_SWIGLU_LIMIT) + return (gate_out * torch.sigmoid(CLAMPED_SWIGLU_ALPHA * gate_out)) * (up_out + 1) + + +# --------------------------------------------------------------------------- +# Registry & dispatch +# --------------------------------------------------------------------------- + +MOE_ACTIVATIONS: Dict[str, Callable[[torch.Tensor, torch.Tensor], torch.Tensor]] = { + "silu": silu_swiglu, + "gelu_tanh": gelu_tanh_glu, + "clamped_swiglu": clamped_swiglu, +} + +SUPPORTED_HIDDEN_ACTS: frozenset[str] = frozenset(MOE_ACTIVATIONS.keys()) + + +def normalize_hidden_act(hidden_act: str | None) -> str: + """Normalize a HF-style ``hidden_act`` string to a canonical MoE act kind.""" + if hidden_act is None or hidden_act == "silu": + return "silu" + if hidden_act in ("gelu_tanh", "gelu_pytorch_tanh"): + return "gelu_tanh" + if hidden_act == "clamped_swiglu": + return "clamped_swiglu" + raise ValueError(f"Unsupported hidden_act={hidden_act!r}. Supported: {sorted(SUPPORTED_HIDDEN_ACTS)}") + + +def check_hidden_act_supported(hidden_act: str, backend: str, supported: frozenset[str]) -> None: + """Raise if ``hidden_act`` is not in the backend's supported set.""" + if hidden_act not in supported: + raise ValueError( + f"MoE backend {backend!r} does not support hidden_act={hidden_act!r}. Supported: {sorted(supported)}" + ) + + +def apply_moe_activation( + hidden_act: str, + gate_out: torch.Tensor, + up_out: torch.Tensor, +) -> torch.Tensor: + """Apply the activation named by ``hidden_act`` to split gate/up tensors. + + Uses explicit ``if`` chain rather than a dict lookup so ``torch.compile`` + specializes on the string value. + """ + if hidden_act == "silu": + return silu_swiglu(gate_out, up_out) + if hidden_act == "gelu_tanh": + return gelu_tanh_glu(gate_out, up_out) + if hidden_act == "clamped_swiglu": + return clamped_swiglu(gate_out, up_out) + raise ValueError(f"Unknown hidden_act={hidden_act!r}. Supported: {sorted(SUPPORTED_HIDDEN_ACTS)}") diff --git a/src/xorl/ops/moe/quack.py b/src/xorl/ops/moe/quack.py index 9237c14e..27742e33 100644 --- a/src/xorl/ops/moe/quack.py +++ b/src/xorl/ops/moe/quack.py @@ -6,6 +6,8 @@ from xorl.distributed.parallel_state import get_parallel_state from xorl.ops.group_gemm.kernel.moe import expert_histogram, moe_gather, moe_index_compute, moe_scatter from xorl.ops.group_gemm.kernel.quack import cumsum_to_cu_seqlens, quack_group_gemm_same_mn, quack_group_gemm_same_nk +from xorl.ops.moe.activations import check_hidden_act_supported +from xorl.ops.moe.triton import _moe_gate_activation, _moe_gate_activation_backward def _debug_ep_enabled() -> bool: @@ -25,52 +27,64 @@ def _scatter_and_cumsum(hidden_states: torch.Tensor, expert_index: torch.Tensor, class QuackMoeExpertsFunction(torch.autograd.Function): - """Memory-optimized: separate gate/up GEMMs, recompute cheap intermediates, - explicit del for dead tensors, in-place add for dgrad.""" + """Fused gate+up GEMM MoE compute. Mirrors ``TritonMoeExpertsFunction``.""" + + SUPPORTED_HIDDEN_ACTS = frozenset({"silu", "gelu_tanh"}) @staticmethod def forward( - ctx, num_experts, gate_weights, expert_index, hidden_states, gate_proj, up_proj, down_proj, gate_up_weight=None + ctx, + num_experts, + gate_weights, + expert_index, + hidden_states, + gate_proj, + up_proj, + down_proj, + gate_up_proj=None, + hidden_act="silu", ): + check_hidden_act_supported(hidden_act, "quack", QuackMoeExpertsFunction.SUPPORTED_HIDDEN_ACTS) + assert gate_up_proj is not None, "QuackMoeExpertsFunction requires a fused gate_up_proj" + ctx.hidden_act = hidden_act + num_tokens = hidden_states.shape[0] + top_k = expert_index.shape[1] + 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 + gate_up_output = quack_group_gemm_same_nk( + a=scatter_output, b=gate_up_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 + I = gate_up_output.shape[-1] // 2 + gate_output = gate_up_output[..., :I] + up_output = gate_up_output[..., I:] - gate_activation = torch.ops.aten.silu(gate_output) + gate_activation = _moe_gate_activation(gate_output, getattr(ctx, "hidden_act", "silu")) 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 projection (NO routing weights inside GEMM — apply after) 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 + a=gated_activation, 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 + del gated_activation + + # Unsort, apply routing weights, reshape+sum (deterministic accumulation) + per_slot = down_output[scatter_index.flatten()].reshape(num_tokens, top_k, -1) + output = (per_slot * gate_weights.unsqueeze(-1)).sum(dim=1) + del down_output, per_slot 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_up_proj, ) return output @@ -78,34 +92,36 @@ def forward( 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_up_proj, ) = ctx.saved_tensors + # Recompute scattered routing weights for backward + 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 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 cheap intermediates (avoids saving them) + # Recompute cheap intermediates scatter_output = moe_scatter(hidden_states, scatter_index) - gate_activation = torch.ops.aten.silu(gate_output) + gate_activation = _moe_gate_activation(gate_output, getattr(ctx, "hidden_act", "silu")) gated_activation = gate_activation * up_output gated_weighted = gated_activation * scattered_gate_weight grad_down_output = moe_scatter(grad_output, scatter_index) - # dgrad FC2 + # FC2 dgrad 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 + # FC2 wgrad grad_down_proj = None if down_proj.requires_grad: grad_down_proj = torch.empty_like(down_proj) @@ -131,170 +147,187 @@ 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 = _moe_gate_activation_backward( + grad_gate_activation, gate_output, getattr(ctx, "hidden_act", "silu") + ) del grad_gate_activation, gate_output - # dgrad FC1: in-place add + # FC1 dgrad + wgrad — fused via gate_up_proj + grad_gate_up_act = torch.cat([grad_gate_output, grad_up_output], dim=-1) + del grad_gate_output, grad_up_output 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 + a=grad_gate_up_act, + b=gate_up_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) + grad_gate_up_proj = None + if gate_up_proj.requires_grad: + grad_gate_up_proj = torch.empty_like(gate_up_proj) quack_group_gemm_same_mn( a=scatter_output, - b=grad_up_output, - c=grad_up_proj, + b=grad_gate_up_act, + c=grad_gate_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 + del grad_gate_up_act, 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 + # Unsort grad + reshape+sum (deterministic, matching forward) + grad_hidden_states = ( + grad_scatter_output[scatter_index.flatten()] + .reshape(hidden_states.shape[0], scatter_index.shape[1], -1) + .sum(dim=1) + ) + + return ( + None, # num_experts + grad_gate_weight, # gate_weights + None, # expert_index + grad_hidden_states, # hidden_states + None, # gate_proj (unused — fused into gate_up_proj) + None, # up_proj (unused — fused into gate_up_proj) + grad_down_proj, # down_proj + grad_gate_up_proj, # gate_up_proj + None, # hidden_act + ) class QuackEPGroupGemm(torch.autograd.Function): """Memory-optimized EP expert GEMM. Recomputes cheap intermediates, explicit del.""" + SUPPORTED_HIDDEN_ACTS = frozenset({"silu", "gelu_tanh"}) + @staticmethod - def forward(ctx, permute_tokens, cumsum, gate_proj, up_proj, down_proj, expert_scores=None): + def forward( + ctx, permute_tokens, cumsum, gate_up_proj, down_proj, intermediate_size, expert_scores=None, hidden_act="silu" + ): + check_hidden_act_supported(hidden_act, "quack", QuackEPGroupGemm.SUPPORTED_HIDDEN_ACTS) + ctx.hidden_act = hidden_act max_M = permute_tokens.shape[0] + I = intermediate_size 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, - expert_scores, - max_M, - cu_seqlens, + ctx, permute_tokens, cumsum, gate_up_proj, down_proj, I, expert_scores, max_M, cu_seqlens ) - 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_up_output = quack_group_gemm_same_nk( + a=permute_tokens, b=gate_up_proj, cumsum_M=cumsum, max_M=max_M, transpose_b=False, cu_seqlens_m=cu_seqlens ) + gate_output = gate_up_output[..., :I] + up_output = gate_up_output[..., I:] - gate_activation = torch.ops.aten.silu(gate_output) + gate_activation = _moe_gate_activation(gate_output, getattr(ctx, "hidden_act", "silu")) 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 projection (NO expert_scores inside — apply after) 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 + if expert_scores is not None: + down_output = down_output * expert_scores.to(down_output.dtype).unsqueeze(-1) + if expert_scores is None: expert_scores = permute_tokens.new_ones(permute_tokens.shape[0]) - ctx.save_for_backward( - permute_tokens, cumsum, gate_proj, up_proj, down_proj, gate_output, up_output, expert_scores - ) + ctx.save_for_backward(permute_tokens, cumsum, gate_up_proj, down_proj, gate_up_output, expert_scores) + ctx.intermediate_size = I return down_output @staticmethod - def _forward_debug(ctx, permute_tokens, cumsum, gate_proj, up_proj, down_proj, expert_scores, max_M, cu_seqlens): + def _forward_debug(ctx, permute_tokens, cumsum, gate_up_proj, down_proj, I, 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)] + ev = [torch.cuda.Event(enable_timing=True) for _ in range(6)] ctx.has_expert_scores = expert_scores is not None ev[0].record() - 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 + gate_up_output = quack_group_gemm_same_nk( + a=permute_tokens, b=gate_up_proj, cumsum_M=cumsum, max_M=max_M, transpose_b=False, cu_seqlens_m=cu_seqlens ) ev[1].record() + gate_output = gate_up_output[..., :I] + up_output = gate_up_output[..., I:] - 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 - ) - ev[2].record() - - gate_activation = torch.ops.aten.silu(gate_output) + gate_activation = _moe_gate_activation(gate_output, getattr(ctx, "hidden_act", "silu")) 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() + ev[2].record() + # Down projection (NO expert_scores inside — apply after, matching normal path) 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 ) - ev[4].record() + ev[3].record() del gated_output + if expert_scores is not None: + down_output = down_output * expert_scores.to(down_output.dtype).unsqueeze(-1) + torch.cuda.synchronize() - t_gate = ev[0].elapsed_time(ev[1]) - t_up = ev[1].elapsed_time(ev[2]) - t_act = ev[2].elapsed_time(ev[3]) - t_down = ev[3].elapsed_time(ev[4]) + t_gate_up = ev[0].elapsed_time(ev[1]) + t_act = ev[1].elapsed_time(ev[2]) + t_down = ev[2].elapsed_time(ev[3]) print( - f"[QuackEP r{rank}] total_M={max_M} G={gate_proj.shape[0]} " - f"K={gate_proj.shape[1]} N_gate={gate_proj.shape[2]} N_down={down_proj.shape[2]}\n" + f"[QuackEP r{rank}] total_M={max_M} G={gate_up_proj.shape[0]} " + f"K={gate_up_proj.shape[1]} N_gate_up={gate_up_proj.shape[2]} N_down={down_proj.shape[2]}\n" f" cu_seqlens: dtype={cu_seqlens.dtype}, len={cu_seqlens.shape[0]}\n" f" permute_tokens: stride={permute_tokens.stride()}, contiguous={permute_tokens.is_contiguous()}\n" - f" gate GEMM: {t_gate:7.2f} ms\n" - f" up GEMM: {t_up:7.2f} ms\n" - f" silu+mul: {t_act:7.2f} ms\n" - f" down GEMM: {t_down:7.2f} ms\n" - f" total: {t_gate + t_up + t_act + t_down:7.2f} ms", + f" gate_up GEMM: {t_gate_up:7.2f} ms\n" + f" silu+mul: {t_act:7.2f} ms\n" + f" down GEMM: {t_down:7.2f} ms\n" + f" total: {t_gate_up + t_act + t_down:7.2f} ms", flush=True, ) if expert_scores is None: expert_scores = permute_tokens.new_ones(permute_tokens.shape[0]) - ctx.save_for_backward( - permute_tokens, cumsum, gate_proj, up_proj, down_proj, gate_output, up_output, expert_scores - ) + ctx.save_for_backward(permute_tokens, cumsum, gate_up_proj, down_proj, gate_up_output, expert_scores) + ctx.intermediate_size = I return down_output @staticmethod def backward(ctx, grad_output): - permute_tokens, cumsum, gate_proj, up_proj, down_proj, gate_output, up_output, expert_scores = ctx.saved_tensors + permute_tokens, cumsum, gate_up_proj, down_proj, gate_up_output, expert_scores = ctx.saved_tensors + I = ctx.intermediate_size max_M = grad_output.shape[0] cu_seqlens_m = cumsum_to_cu_seqlens(cumsum) - # Recompute cheap intermediates - gate_activation = torch.ops.aten.silu(gate_output) + gate_output = gate_up_output[..., :I] + up_output = gate_up_output[..., I:] + + gate_activation = _moe_gate_activation(gate_output, getattr(ctx, "hidden_act", "silu")) 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) + + # Forward was: out = down_GEMM(gated_output) * expert_scores + # Skip the extra down-GEMM when expert_scores doesn't require a gradient + # (e.g., train_router=False causes routing_weights to be detached upstream, + # so ctx.needs_input_grad[5] is False and grad_expert_scores would be unused). + grad_expert_scores = None + if ctx.has_expert_scores and ctx.needs_input_grad[5]: + 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_m + ) + grad_expert_scores = (down_output * grad_output).sum(dim=-1).to(expert_scores_dtype) + del down_output + + grad_scaled = grad_output * expert_scores.unsqueeze(-1) # dgrad FC2 - grad_gated_weighted = quack_group_gemm_same_nk( - a=grad_output, b=down_proj, cumsum_M=cumsum, max_M=max_M, transpose_b=True, cu_seqlens_m=cu_seqlens_m + grad_gated_output = quack_group_gemm_same_nk( + a=grad_scaled, b=down_proj, cumsum_M=cumsum, max_M=max_M, transpose_b=True, cu_seqlens_m=cu_seqlens_m ) # wgrad FC2 @@ -302,76 +335,84 @@ 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_output, + a=gated_output, + b=grad_scaled, c=grad_down_proj, cumsum_K=cumsum, max_K=max_M, transpose_a=True, - transpose_b=False, cu_seqlens_k=cu_seqlens_m, ) - grad_expert_scores = None - if ctx.has_expert_scores: - grad_expert_scores = (grad_gated_weighted * gated_output).sum(dim=-1).to(expert_scores_dtype) - del gated_output, gated_weighted - - grad_gated_output = grad_gated_weighted * expert_scores.unsqueeze(-1) - del grad_gated_weighted + del gated_output, grad_scaled # 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_gated_output, gate_activation, up_output, gate_up_output + grad_gate_output = _moe_gate_activation_backward( + grad_gate_activation, gate_output, getattr(ctx, "hidden_act", "silu") + ) del grad_gate_activation, gate_output - # dgrad FC1: in-place add + # Fused dgrad FC1 + grad_gate_up_act = torch.cat([grad_gate_output, grad_up_output], dim=-1) + del grad_gate_output, grad_up_output 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 + a=grad_gate_up_act, + b=gate_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) + # Fused wgrad FC1 + grad_gate_up_proj = None + if gate_up_proj.requires_grad: + grad_gate_up_proj = torch.empty_like(gate_up_proj) quack_group_gemm_same_mn( a=permute_tokens, - b=grad_up_output, - c=grad_up_proj, + b=grad_gate_up_act, + c=grad_gate_up_proj, cumsum_K=cumsum, max_K=max_M, transpose_a=True, transpose_b=False, cu_seqlens_k=cu_seqlens_m, ) - del grad_up_output - - return grad_permute_tokens, None, grad_gate_proj, grad_up_proj, grad_down_proj, grad_expert_scores + del grad_gate_up_act + + return ( + grad_permute_tokens, + None, # cumsum + grad_gate_up_proj, + grad_down_proj, + None, # intermediate_size + grad_expert_scores, + None, # hidden_act + ) class QuackTPMoeExpertsFunction(torch.autograd.Function): """Memory-optimized TP expert function. Recomputes cheap intermediates, explicit del + all-reduce.""" + SUPPORTED_HIDDEN_ACTS = frozenset({"silu", "gelu_tanh"}) + @staticmethod - def forward(ctx, num_experts, gate_weights, expert_index, hidden_states, gate_proj, up_proj, down_proj, tp_group): + def forward( + ctx, + num_experts, + gate_weights, + expert_index, + hidden_states, + gate_proj, + up_proj, + down_proj, + tp_group, + hidden_act="silu", + ): + check_hidden_act_supported(hidden_act, "quack", QuackTPMoeExpertsFunction.SUPPORTED_HIDDEN_ACTS) + ctx.hidden_act = hidden_act 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) @@ -384,7 +425,7 @@ def forward(ctx, num_experts, gate_weights, expert_index, hidden_states, gate_pr ) del scatter_output - gate_activation = torch.ops.aten.silu(gate_output) + gate_activation = _moe_gate_activation(gate_output, getattr(ctx, "hidden_act", "silu")) gated_activation = gate_activation * up_output del gate_activation @@ -437,7 +478,7 @@ def backward(ctx, grad_output): # Recompute cheap intermediates (avoids saving them) scatter_output = moe_scatter(hidden_states, scatter_index) - gate_activation = torch.ops.aten.silu(gate_output) + gate_activation = _moe_gate_activation(gate_output, getattr(ctx, "hidden_act", "silu")) gated_activation = gate_activation * up_output gated_weighted = gated_activation * scattered_gate_weight @@ -474,7 +515,9 @@ 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 = _moe_gate_activation_backward( + grad_gate_activation, gate_output, getattr(ctx, "hidden_act", "silu") + ) del grad_gate_activation, gate_output # dgrad FC1: in-place add @@ -518,7 +561,17 @@ def backward(ctx, grad_output): handle.wait() 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 + return ( + None, + grad_gate_weight, + None, + grad_hidden_states, + grad_gate_proj, + grad_up_proj, + grad_down_proj, + None, + None, + ) # hidden_act def quack_moe_forward( @@ -530,8 +583,8 @@ def quack_moe_forward( gate_proj: torch.Tensor, up_proj: torch.Tensor, down_proj: torch.Tensor, - gate_up_weight: torch.Tensor = None, - **kwargs, + gate_up_proj: torch.Tensor = None, + hidden_act: str = "silu", ): """Forward pass for MoE experts using quack group GEMM (local/TP only). @@ -543,9 +596,25 @@ def quack_moe_forward( if parallel_state.tp_enabled: tp_group = parallel_state.tp_mesh.get_group() return QuackTPMoeExpertsFunction.apply( - num_experts, routing_weights, selected_experts, hidden_states, gate_proj, up_proj, down_proj, tp_group + num_experts, + routing_weights, + selected_experts, + hidden_states, + gate_proj, + up_proj, + down_proj, + tp_group, + hidden_act, ) return QuackMoeExpertsFunction.apply( - num_experts, routing_weights, selected_experts, hidden_states, gate_proj, up_proj, down_proj, gate_up_weight + num_experts, + routing_weights, + selected_experts, + hidden_states, + gate_proj, + up_proj, + down_proj, + gate_up_proj, + hidden_act, ) diff --git a/src/xorl/ops/moe/triton.py b/src/xorl/ops/moe/triton.py index d4324bd2..8a7ff5ce 100644 --- a/src/xorl/ops/moe/triton.py +++ b/src/xorl/ops/moe/triton.py @@ -1,6 +1,13 @@ import torch - -from xorl.ops.fused_silu_and_mul import silu_and_mul, silu_and_mul_backward +import torch.nn.functional as F + +# Re-export canonical activation utilities so callers that historically +# imported from ``xorl.ops.moe.triton`` keep working. +from xorl.ops.moe.activations import ( # noqa: F401 + SUPPORTED_HIDDEN_ACTS, + check_hidden_act_supported, + normalize_hidden_act, +) from xorl.utils.import_utils import is_fused_moe_available @@ -8,26 +15,48 @@ 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_gather, moe_index_compute, moe_scatter, ) +def _moe_gate_activation(gate_output: torch.Tensor, hidden_act: str = "silu") -> torch.Tensor: + """Apply gate activation by kind.""" + if hidden_act == "gelu_tanh": + return F.gelu(gate_output, approximate="tanh") + return torch.ops.aten.silu(gate_output) + + +def _moe_gate_activation_backward( + grad: torch.Tensor, gate_output: torch.Tensor, hidden_act: str = "silu" +) -> torch.Tensor: + """Backward for gate activation.""" + if hidden_act == "gelu_tanh": + with torch.enable_grad(): + g = gate_output.detach().requires_grad_(True) + a = F.gelu(g, approximate="tanh") + return torch.autograd.grad(a, g, grad)[0] + return torch.ops.aten.silu_backward(grad, gate_output) + + class TritonEPGroupGemm(torch.autograd.Function): """EP expert MLP with fused gate+up GEMM. Zero-copy weight references. - Forward: single ``x @ gate_up_proj`` GEMM → split → SwiGLU → down GEMM. + Forward: single ``x @ gate_up_proj`` GEMM → split → GLU activation → 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. """ + SUPPORTED_HIDDEN_ACTS = frozenset({"silu", "gelu_tanh"}) + @staticmethod - def forward(ctx, permute_tokens, cumsum, gate_up_proj, down_proj, intermediate_size, expert_scores=None): + def forward( + ctx, permute_tokens, cumsum, gate_up_proj, down_proj, intermediate_size, expert_scores=None, hidden_act="silu" + ): + check_hidden_act_supported(hidden_act, "triton", TritonEPGroupGemm.SUPPORTED_HIDDEN_ACTS) max_M = permute_tokens.shape[0] I = intermediate_size ctx.has_expert_scores = expert_scores is not None + ctx.hidden_act = hidden_act gate_up_output = group_gemm_same_nk( a=permute_tokens, @@ -37,22 +66,25 @@ def forward(ctx, permute_tokens, cumsum, gate_up_proj, down_proj, intermediate_s transpose_a=False, transpose_b=False, ) + gate_output = gate_up_output[..., :I] + up_output = gate_up_output[..., I:] - # 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) + gate_activation = _moe_gate_activation(gate_output, getattr(ctx, "hidden_act", "silu")) + gated_output = gate_activation * up_output + del gate_activation + # Down projection (NO expert_scores inside GEMM — apply after) down_output = group_gemm_same_nk( a=gated_output, b=down_proj, cumsum_M=cumsum, max_M=max_M, - transpose_a=False, - transpose_b=False, ) del gated_output + if expert_scores is not None: + down_output = down_output * expert_scores.to(down_output.dtype).unsqueeze(-1) + 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) @@ -63,16 +95,32 @@ def forward(ctx, permute_tokens, cumsum, gate_up_proj, down_proj, intermediate_s @staticmethod def backward(ctx, grad_output): permute_tokens, cumsum, gate_up_proj, down_proj, gate_up_output, expert_scores = ctx.saved_tensors + I = ctx.intermediate_size max_M = grad_output.shape[0] - gated_output = silu_and_mul(gate_up_output) + gate_output = gate_up_output[..., :I] + up_output = gate_up_output[..., I:] + + gate_activation = _moe_gate_activation(gate_output, getattr(ctx, "hidden_act", "silu")) + 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) + + # Forward was: out = down_GEMM(gated_output) * expert_scores + # Skip the extra down-GEMM when expert_scores doesn't require a gradient + # (e.g., train_router=False causes routing_weights to be detached upstream, + # so ctx.needs_input_grad[5] is False and grad_expert_scores would be unused). + grad_expert_scores = None + if ctx.has_expert_scores and ctx.needs_input_grad[5]: + down_output = group_gemm_same_nk(a=gated_output, b=down_proj, cumsum_M=cumsum, max_M=max_M) + grad_expert_scores = (down_output * grad_output).sum(dim=-1).to(expert_scores_dtype) + del down_output + + grad_scaled = grad_output * expert_scores.unsqueeze(-1) # dgrad FC2 - grad_gated_weighted = group_gemm_same_nk( - a=grad_output, + grad_gated_output = group_gemm_same_nk( + a=grad_scaled, b=down_proj, cumsum_M=cumsum, max_M=max_M, @@ -84,26 +132,28 @@ 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_output, + a=gated_output, + b=grad_scaled, c=grad_down_proj, cumsum_K=cumsum, max_K=max_M, transpose_a=True, - transpose_b=False, ) - grad_expert_scores = None - if ctx.has_expert_scores: - grad_expert_scores = (grad_gated_weighted * gated_output).sum(dim=-1).to(expert_scores_dtype) - del gated_output, gated_weighted + del gated_output, grad_scaled - grad_gated_output = grad_gated_weighted * expert_scores.unsqueeze(-1) - del grad_gated_weighted + # Activation backward + grad_up_output = gate_activation * grad_gated_output + grad_gate_activation = grad_gated_output * up_output + del grad_gated_output, gate_activation, up_output, gate_up_output - # 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_gate_output = _moe_gate_activation_backward( + grad_gate_activation, gate_output, getattr(ctx, "hidden_act", "silu") + ) + del grad_gate_activation, gate_output + # Fused dgrad FC1 + grad_gate_up_act = torch.cat([grad_gate_output, grad_up_output], dim=-1) + del grad_gate_output, grad_up_output grad_permute_tokens = group_gemm_same_nk( a=grad_gate_up_act, b=gate_up_proj, @@ -134,6 +184,7 @@ def backward(ctx, grad_output): grad_down_proj, None, # intermediate_size grad_expert_scores, + None, # hidden_act ) @@ -145,6 +196,8 @@ class TritonMoeExpertsFunction(torch.autograd.Function): to free dead tensors immediately. """ + SUPPORTED_HIDDEN_ACTS = frozenset({"silu", "gelu_tanh"}) + @staticmethod def forward( ctx, @@ -155,63 +208,51 @@ def forward( gate_proj, up_proj, down_proj, - gate_up_weight=None, # Ignored (kept for API compat) + gate_up_proj=None, + hidden_act="silu", ): - # Token dispatch: compute histogram, scatter index, and scatter tokens + check_hidden_act_supported(hidden_act, "triton", TritonMoeExpertsFunction.SUPPORTED_HIDDEN_ACTS) + ctx.hidden_act = hidden_act + num_tokens = hidden_states.shape[0] + top_k = expert_index.shape[1] + + # Token dispatch: sort by expert 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] - # 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, - ) - up_output = group_gemm_same_nk( + assert gate_up_proj is not None, "TritonMoeExpertsFunction requires a fused gate_up_proj" + gate_up_output = group_gemm_same_nk( a=scatter_output, - b=up_proj, + b=gate_up_proj, cumsum_M=cumsum_t, max_M=max_M, - transpose_a=False, - transpose_b=False, ) - del scatter_output + I = gate_up_output.shape[-1] // 2 + gate_output = gate_up_output[..., :I] + up_output = gate_up_output[..., I:] - # SiLU activation + element-wise multiply (bf16 like native backend) - gate_activation = torch.ops.aten.silu(gate_output) + # Activation + GLU + gate_activation = _moe_gate_activation(gate_output, getattr(ctx, "hidden_act", "silu")) gated_activation = gate_activation * up_output del gate_activation - # Apply routing weights in scattered layout - 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 projection + # Down projection (NO routing weights inside GEMM — apply after) down_output = group_gemm_same_nk( - a=gated_weighted, + a=gated_activation, b=down_proj, cumsum_M=cumsum_t, max_M=max_M, - transpose_a=False, - transpose_b=False, ) - del gated_weighted + del gated_activation - # Gather and reshape - output = moe_gather(down_output, scatter_index).reshape(hidden_states.shape) - del down_output + # Unsort, apply routing weights, reshape+sum (deterministic accumulation) + per_slot = down_output[scatter_index.flatten()].reshape(num_tokens, top_k, -1) + output = (per_slot * gate_weights.unsqueeze(-1)).sum(dim=1) + del down_output, per_slot - # 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, @@ -222,7 +263,7 @@ def forward( cumsum_t, gate_output, up_output, - scattered_gate_weight, + gate_up_proj, ) return output @@ -239,21 +280,25 @@ def backward(ctx, grad_output): cumsum_t, gate_output, up_output, - scattered_gate_weight, + gate_up_proj, ) = ctx.saved_tensors + # Recompute scattered routing weights for backward + 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 grad_output = grad_output.view(-1, grad_output.shape[-1]) max_M = grad_output.shape[0] - # Recompute cheap intermediates (avoids saving them) + # Recompute cheap intermediates scatter_output = moe_scatter(hidden_states, scatter_index) - gate_activation = torch.ops.aten.silu(gate_output) + gate_activation = _moe_gate_activation(gate_output, getattr(ctx, "hidden_act", "silu")) 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 @ down_proj^T + # FC2 dgrad grad_gated_weighted = group_gemm_same_nk( a=grad_down_output, b=down_proj, @@ -262,7 +307,7 @@ def backward(ctx, grad_output): transpose_b=True, ) - # FC2 wgrad: gated_weighted^T @ grad + # FC2 wgrad grad_down_proj = None if down_proj.requires_grad: grad_down_proj = torch.empty_like(down_proj) @@ -273,7 +318,6 @@ def backward(ctx, grad_output): cumsum_K=cumsum_t, max_K=max_M, transpose_a=True, - transpose_b=False, ) del grad_down_output, gated_weighted @@ -283,69 +327,55 @@ def backward(ctx, grad_output): grad_gate_weight = grad_gate_weight.reshape(gate_weights.shape) del gated_activation, grad_gated_weighted - # Activation backward (separate ops, matching TritonEPGroupGemm pattern) + # 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) + grad_gate_output = _moe_gate_activation_backward( + grad_gate_activation, gate_output, getattr(ctx, "hidden_act", "silu") + ) del grad_gate_activation, gate_output - # FC1 dgrad: separate GEMMs, in-place add (avoids allocating sum tensor) + # FC1 dgrad + wgrad — fused via gate_up_proj + grad_gate_up_act = torch.cat([grad_gate_output, grad_up_output], dim=-1) + del grad_gate_output, grad_up_output 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, + a=grad_gate_up_act, + b=gate_up_proj, cumsum_M=cumsum_t, max_M=max_M, transpose_b=True, ) - - # FC1 wgrad: separate GEMMs (avoids concatenated grad weight alloc + .contiguous() copies) - 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) + 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=scatter_output, - b=grad_up_output, - c=grad_up_proj, + b=grad_gate_up_act, + c=grad_gate_up_proj, cumsum_K=cumsum_t, max_K=max_M, transpose_a=True, - transpose_b=False, ) - del grad_up_output, scatter_output + del grad_gate_up_act, scatter_output - # Gather gradient for hidden_states - grad_hidden_states = moe_gather(grad_scatter_output, scatter_index).reshape(hidden_states.shape) + # Unsort grad + reshape+sum (deterministic, matching forward) + grad_hidden_states = ( + grad_scatter_output[scatter_index.flatten()] + .reshape(hidden_states.shape[0], scatter_index.shape[1], -1) + .sum(dim=1) + ) return ( None, # num_experts grad_gate_weight, # gate_weights None, # expert_index grad_hidden_states, # hidden_states - grad_gate_proj, # gate_proj - grad_up_proj, # up_proj + None, # gate_proj (unused — fused into gate_up_proj) + None, # up_proj (unused — fused into gate_up_proj) grad_down_proj, # down_proj - None, # gate_up_weight + grad_gate_up_proj, # gate_up_proj + None, # hidden_act ) @@ -358,8 +388,8 @@ def triton_moe_forward( gate_proj: torch.Tensor, up_proj: torch.Tensor, down_proj: torch.Tensor, - gate_up_weight: torch.Tensor = None, - **kwargs, + gate_up_proj: torch.Tensor = None, + hidden_act: str = "silu", ): """Forward pass for MoE experts using Triton group GEMM (local, single-GPU). @@ -374,7 +404,8 @@ def triton_moe_forward( gate_proj: Gate projection weights, shape [num_experts, hidden_dim, intermediate_size]. up_proj: Up projection weights, shape [num_experts, hidden_dim, intermediate_size]. down_proj: Down projection weights, shape [num_experts, intermediate_size, hidden_dim]. - gate_up_weight: Optional pre-concatenated weights [num_experts, hidden_dim, 2*intermediate_size]. + gate_up_proj: Pre-fused weights [num_experts, hidden_dim, 2*intermediate_size]. + hidden_act: Activation kind ("silu" or "gelu_tanh"). Returns: Output hidden states, shape [num_tokens, hidden_dim]. @@ -387,5 +418,6 @@ def triton_moe_forward( gate_proj, up_proj, down_proj, - gate_up_weight, + gate_up_proj, + hidden_act, ) diff --git a/src/xorl/ops/quack/cross_entropy.py b/src/xorl/ops/quack/cross_entropy.py index bbc6de34..bddca72d 100644 --- a/src/xorl/ops/quack/cross_entropy.py +++ b/src/xorl/ops/quack/cross_entropy.py @@ -226,7 +226,7 @@ def kernel( copy(tXrdX, tXgdX) -@torch.library.custom_op("quack::cross_entropy_fwd_out", mutates_args={"loss", "lse", "dx"}) +@torch.library.custom_op("xorl_quack::cross_entropy_fwd_out", mutates_args={"loss", "lse", "dx"}) def cross_entropy_fwd_out( x: Tensor, target: Tensor, @@ -524,7 +524,7 @@ def _cross_entropy_backward( _cross_entropy_backward.compile_cache = {} -@torch.library.custom_op("quack::cross_entropy_bwd_out", mutates_args={"dx"}) +@torch.library.custom_op("xorl_quack::cross_entropy_bwd_out", mutates_args={"dx"}) def cross_entropy_bwd_out( x: torch.Tensor, target: torch.Tensor, diff --git a/src/xorl/ops/quack/gemm_interface.py b/src/xorl/ops/quack/gemm_interface.py index fc71fb31..de9af1a4 100644 --- a/src/xorl/ops/quack/gemm_interface.py +++ b/src/xorl/ops/quack/gemm_interface.py @@ -313,7 +313,7 @@ def gemm( @torch.library.custom_op( - "quack::gemm_out", + "xorl_quack::gemm_out", mutates_args=("out",), device_types="cuda", # We have to split out alpha and alpha_tensor since torch.library requires @@ -469,7 +469,7 @@ def gemm_add( @torch.library.custom_op( - "quack::gemm_add_out", + "xorl_quack::gemm_add_out", mutates_args=("out",), device_types="cuda", # We have to split out alpha and alpha_tensor since torch.library requires @@ -627,7 +627,7 @@ def gemm_add_inplace( @torch.library.custom_op( - "quack::gemm_add_inplace", + "xorl_quack::gemm_add_inplace", mutates_args=("out",), device_types="cuda", # We have to split out alpha and alpha_tensor since torch.library requires @@ -720,7 +720,7 @@ def gemm_act( @torch.library.custom_op( - "quack::gemm_act_out", + "xorl_quack::gemm_act_out", mutates_args=("preact_out", "postact_out"), device_types="cuda", schema="(Tensor A, Tensor B, Tensor(a2!)? preact_out, Tensor(a3!) postact_out, Tensor? C=None, Tensor? bias=None, str? activation=None, Tensor? cu_seqlens_m=None, Tensor? A_idx=None, bool dynamic_scheduler=False, bool tuned=True) -> ()", @@ -800,7 +800,7 @@ def gemm_dact( @torch.library.custom_op( - "quack::gemm_dact_out", + "xorl_quack::gemm_dact_out", mutates_args=("dx_out", "postact_out"), device_types="cuda", schema="(Tensor A, Tensor B, Tensor PreAct, Tensor(a3!) dx_out, Tensor(a4!) postact_out, str? activation=None, Tensor? cu_seqlens_m=None, Tensor? A_idx=None, bool dynamic_scheduler=True, bool tuned=True) -> ()", @@ -935,7 +935,7 @@ def gemm_dgated_ref( @torch.library.custom_op( - "quack::gemm_symmetric_out", + "xorl_quack::gemm_symmetric_out", mutates_args=("out",), device_types="cuda", schema="(Tensor A, Tensor B, Tensor(a2!) out, Tensor? C=None, bool dynamic_scheduler=False, float alpha=1.0, float beta=1.0) -> ()", @@ -1217,7 +1217,7 @@ def gemm_gated( @torch.library.custom_op( - "quack::gemm_gated_out", + "xorl_quack::gemm_gated_out", mutates_args=("preact_out", "postact_out"), device_types="cuda", schema="(Tensor A, Tensor B, Tensor(a2!)? preact_out, Tensor(a3!) postact_out, Tensor? C=None, Tensor? bias=None, str activation='swiglu', Tensor? cu_seqlens_m=None, Tensor? A_idx=None, bool dynamic_scheduler=False, bool tuned=True) -> ()", @@ -1294,7 +1294,7 @@ def gemm_dgated( @torch.library.custom_op( - "quack::gemm_dgated_out", + "xorl_quack::gemm_dgated_out", mutates_args=("dx_out", "postact_out"), device_types="cuda", schema="(Tensor A, Tensor B, Tensor PreAct, Tensor(a3!) dx_out, Tensor(a4!) postact_out, Tensor? colvec_scale=None, str activation='swiglu', bool colvec_reduce=False, Tensor? cu_seqlens_m=None, Tensor? A_idx=None, bool dynamic_scheduler=True, bool tuned=True) -> Tensor?", @@ -1330,7 +1330,7 @@ def gemm_dgated_out( ) -@torch.library.register_fake("quack::gemm_dgated_out") +@torch.library.register_fake("xorl_quack::gemm_dgated_out") def gemm_dgated_out_fake( A: Tensor, # (M, K) or (L, M, K) or (total_M, K) if varlen_m or (whatever, K) if gather_A with varlen_m B: Tensor, # (K, N) or (L, K, N) @@ -1362,8 +1362,8 @@ def gemm_dgated_out_fake( # try: # from torch._inductor.fx_passes.reinplace import InplaceableOp # torch._inductor.fx_passes.reinplace.inplaceable_ops.update({ -# torch.ops.quack.gemm_add_out.default: -# InplaceableOp(torch.ops.quack.gemm_add_inplace.default, mutated_arg=2) +# torch.ops.xorl_quack.gemm_add_out.default: +# InplaceableOp(torch.ops.xorl_quack.gemm_add_inplace.default, mutated_arg=2) # }) # except ImportError: # pass diff --git a/src/xorl/ops/quack/rmsnorm.py b/src/xorl/ops/quack/rmsnorm.py index 43f239f9..e8c8ed62 100644 --- a/src/xorl/ops/quack/rmsnorm.py +++ b/src/xorl/ops/quack/rmsnorm.py @@ -267,7 +267,7 @@ def kernel( @torch.library.custom_op( - "quack::_rmsnorm_fwd", + "xorl_quack::_rmsnorm_fwd", mutates_args=("out", "rstd", "mean", "residual_out"), device_types="cuda", # We need to specify the schema manually since we're mutating an optional tensor @@ -747,7 +747,7 @@ def _get_sm_count(N: int, device: torch.device) -> int: @torch.library.custom_op( - "quack::_rmsnorm_bwd", + "xorl_quack::_rmsnorm_bwd", mutates_args={"dx", "dw_partial", "db_partial", "dresidual"}, device_types="cuda", # We need to specify the schema manually since we're mutating an optional tensor diff --git a/src/xorl/ops/quack/softmax.py b/src/xorl/ops/quack/softmax.py index 18bcf849..65f89af8 100644 --- a/src/xorl/ops/quack/softmax.py +++ b/src/xorl/ops/quack/softmax.py @@ -152,7 +152,7 @@ def kernel( copy(tXrO, tXgO) -@torch.library.custom_op("quack::_softmax_fwd", mutates_args={"out"}) +@torch.library.custom_op("xorl_quack::_softmax_fwd", mutates_args={"out"}) def _softmax_fwd(x: torch.Tensor, out: torch.Tensor) -> None: """Softmax forward pass. Args: @@ -312,7 +312,7 @@ def kernel( copy(tdXrdX, tdXgdX) -@torch.library.custom_op("quack::_softmax_backward", mutates_args={"dx"}) +@torch.library.custom_op("xorl_quack::_softmax_backward", mutates_args={"dx"}) def _softmax_backward(dy: torch.Tensor, y: torch.Tensor, dx: torch.Tensor) -> None: """Softmax backward pass. Args: diff --git a/src/xorl/ops/quack/topk.py b/src/xorl/ops/quack/topk.py index 40c19d69..611e1e53 100644 --- a/src/xorl/ops/quack/topk.py +++ b/src/xorl/ops/quack/topk.py @@ -205,7 +205,7 @@ def kernel( cute.autovec_copy(topk_indices[None, i], mIndices_store[None, col]) -@torch.library.custom_op("quack::_topk_fwd", mutates_args={"values", "indices"}) +@torch.library.custom_op("xorl_quack::_topk_fwd", mutates_args={"values", "indices"}) def _topk_fwd(x: torch.Tensor, k: int, softmax: bool, values: torch.Tensor, indices: torch.Tensor) -> None: """Top-k forward pass. Args: @@ -421,7 +421,7 @@ def kernel( copy_dx(tXrdX, tXgdX) -@torch.library.custom_op("quack::_topk_bwd", mutates_args={"dx"}) +@torch.library.custom_op("xorl_quack::_topk_bwd", mutates_args={"dx"}) def _topk_bwd( dvalues: torch.Tensor, values: Optional[torch.Tensor], diff --git a/src/xorl/optim/__init__.py b/src/xorl/optim/__init__.py index cef34cee..c7e77d47 100644 --- a/src/xorl/optim/__init__.py +++ b/src/xorl/optim/__init__.py @@ -1,4 +1,5 @@ from .anyprecision_adamw import AnyPrecisionAdamW +from .distsignsgd import DistSignSGD from .lr_scheduler import build_lr_scheduler from .multi_optimizer import MultiOptimizer from .muon import Muon @@ -10,6 +11,7 @@ "AnyPrecisionAdamW", "build_lr_scheduler", "build_optimizer", + "DistSignSGD", "MultiOptimizer", "Muon", "SignSGD", diff --git a/src/xorl/optim/anyprecision_adamw.py b/src/xorl/optim/anyprecision_adamw.py index 3785fbc0..97e2dbbc 100644 --- a/src/xorl/optim/anyprecision_adamw.py +++ b/src/xorl/optim/anyprecision_adamw.py @@ -1,6 +1,8 @@ import torch from torch.optim.optimizer import Optimizer +from .cautious import apply_cautious_decay_ + class AnyPrecisionAdamW(Optimizer): def __init__( @@ -14,6 +16,7 @@ def __init__( momentum_dtype=torch.bfloat16, variance_dtype=torch.bfloat16, compensation_buffer_dtype=torch.bfloat16, + cautious=False, ): defaults = { "lr": lr, @@ -24,6 +27,7 @@ def __init__( "momentum_dtype": momentum_dtype, "variance_dtype": variance_dtype, "compensation_buffer_dtype": compensation_buffer_dtype, + "cautious": cautious, } super().__init__(params, defaults) @@ -32,6 +36,11 @@ def step(self, closure=None): """ Performs a single optimization step. + When ``cautious=True``, decoupled weight decay is masked by + ``I(exp_avg * param >= 0)`` per Chen et al. "Cautious Weight Decay" + (arXiv:2510.12402). ``sign(exp_avg)`` matches ``sign(u_t)`` since the + Adam preconditioner denominator is strictly positive. + Args: closure (callable, optional): A closure that reevaluates the model and returns the loss. """ @@ -46,6 +55,7 @@ def step(self, closure=None): weight_decay = group["weight_decay"] eps = group["eps"] use_kahan_summation = group["use_kahan_summation"] + cautious = group.get("cautious", False) momentum_dtype = group["momentum_dtype"] variance_dtype = group["variance_dtype"] @@ -73,12 +83,21 @@ def step(self, closure=None): 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) + # Cautious weight decay must use the post-update first moment + # (its sign matches the optimizer update direction). Apply + # decay AFTER moments are updated and BEFORE the parameter + # update, against the pre-update parameter values. + apply_cautious_decay_( + p.data, + update_sign_proxy=exp_avg, + lr=lr, + weight_decay=weight_decay, + cautious=cautious, + ) + bias_correction1 = 1 - beta1**step step_size = lr / bias_correction1 diff --git a/src/xorl/optim/cautious.py b/src/xorl/optim/cautious.py new file mode 100644 index 00000000..096dbf27 --- /dev/null +++ b/src/xorl/optim/cautious.py @@ -0,0 +1,47 @@ +"""Cautious Weight Decay (CWD) primitives. + +CWD applies decoupled weight decay only along coordinates where the optimizer +update and the parameter share a sign:: + + x_{t+1} = x_t - eta * (u_t + lambda * I(u_t * x_t >= 0) * x_t) + +Compared to standard decoupled decay (``x_{t+1} = (1 - eta*lambda) * x_t - eta * u_t``), +the decay term is masked elementwise by ``I(u_t * x_t >= 0)``. The mask uses +the sign of the *optimizer update* ``u_t`` (after preconditioning / +orthogonalization), not the raw gradient. + +Reference: Chen et al., "Cautious Weight Decay" (arXiv:2510.12402). +""" + +import torch + + +def apply_cautious_decay_( + param: torch.Tensor, + update_sign_proxy: torch.Tensor, + *, + lr: float, + weight_decay: float, + cautious: bool, +) -> None: + """In-place decoupled weight decay, optionally masked by ``I(u * x >= 0)``. + + When ``cautious=False`` this is the standard ``param *= 1 - lr * weight_decay``. + When ``cautious=True`` the decay factor becomes + ``1 - lr * weight_decay * I(update_sign_proxy * param >= 0)`` elementwise. + + ``update_sign_proxy`` only needs to share its sign with the optimizer + update direction ``u_t``. For Adam-family optimizers ``exp_avg`` is a valid + proxy (the preconditioner denominator is strictly positive). For SignSGD + ``grad`` is a valid proxy. For Muon the post-Newton-Schulz ``update`` + tensor is the proxy (it *is* ``u_t``). + + No-op when ``weight_decay == 0``. + """ + if weight_decay == 0.0: + return + if not cautious: + param.mul_(1.0 - lr * weight_decay) + return + mask = (update_sign_proxy * param >= 0).to(param.dtype) + param.mul_(mask.mul_(-lr * weight_decay).add_(1.0)) diff --git a/src/xorl/optim/distsignsgd.py b/src/xorl/optim/distsignsgd.py new file mode 100644 index 00000000..451e1ba4 --- /dev/null +++ b/src/xorl/optim/distsignsgd.py @@ -0,0 +1,270 @@ +from __future__ import annotations + +from collections.abc import Sequence +from typing import Optional + +import torch +import torch.distributed as dist +from torch.distributed.fsdp import FSDPModule +from torch.optim.optimizer import Optimizer + +from ..distributed.parallel_state import get_parallel_state + + +try: + from torch.distributed._tensor import DTensor +except ImportError: # pragma: no cover - torch 2.10+ always provides DTensor here + DTensor = None + + +class _DefaultReduceScatterComm: + """Minimal reduce-scatter implementation matching FSDP2's comm interface.""" + + def allocate( + self, + size: Sequence[int | torch.SymInt], + *, + dtype: torch.dtype, + device: torch.device, + ) -> torch.Tensor: + return torch.empty(*size, dtype=dtype, device=device) + + def __call__( + self, + output_tensor: torch.Tensor, + input_tensor: torch.Tensor, + group: dist.ProcessGroup, + op: dist.ReduceOp | dist.ReduceOp.RedOpType, + async_op: bool = False, + ) -> Optional[dist.Work]: + return dist.reduce_scatter_tensor( + output=output_tensor, + input=input_tensor, + group=group, + op=op, + async_op=async_op, + ) + + +class DistSignReduceScatter: + """Signs the flattened gradient buffer before FSDP2 reduces it.""" + + def __init__(self, inner_comm: Optional[object] = None, *, sp_group: Optional[dist.ProcessGroup] = None): + self._inner = inner_comm or _DefaultReduceScatterComm() + self._sp_group = sp_group + + def allocate( + self, + size: Sequence[int | torch.SymInt], + *, + dtype: torch.dtype, + device: torch.device, + ) -> torch.Tensor: + return self._inner.allocate(size, dtype=dtype, device=device) + + def __call__( + self, + output_tensor: torch.Tensor, + input_tensor: torch.Tensor, + group: dist.ProcessGroup, + op: dist.ReduceOp | dist.ReduceOp.RedOpType, + async_op: bool = False, + ) -> Optional[dist.Work]: + if input_tensor.is_sparse: + raise RuntimeError("DistSignSGD does not support sparse gradients.") + if self._sp_group is not None: + # Exact SP sums must happen before the nonlinear sign. + dist.all_reduce(input_tensor, op=dist.ReduceOp.SUM, group=self._sp_group) + input_tensor.sign_() + # Force SUM regardless of `op`. FSDP2 may pass AVG (e.g. with + # reduce_dtype=fp32), which would divide the sign-vote sum by N — and + # the trainer's scale factor already divides by the actual voter total. + # Inheriting AVG would double-divide and silently shrink updates. + return self._inner( + output_tensor=output_tensor, + input_tensor=input_tensor, + group=group, + op=dist.ReduceOp.SUM, + async_op=async_op, + ) + + +def _local_sign_hook(grad: torch.Tensor) -> torch.Tensor: + if grad.is_sparse: + raise RuntimeError("DistSignSGD does not support sparse gradients.") + return torch.sign(grad) + + +_FSDP_INTERNALS_ERROR = ( + "DistSignSGD relies on private torch.distributed.fsdp internals " + "(`FSDPModule._get_fsdp_state()._fsdp_param_group.fsdp_params[*].sharded_param`). " + "These have changed in your torch build, so the FSDP-managed vs local-hook split " + "cannot be computed safely. Pin a compatible torch version or update DistSignSGD." +) + + +def _iter_fsdp_managed_sharded_params(model: torch.nn.Module): + """Yield ``sharded_param`` tensors for every FSDP-managed parameter in ``model``. + + Raises a clear RuntimeError if any of the unstable internal attributes + (`_get_fsdp_state`, `_fsdp_param_group`, `fsdp_params`, `sharded_param`) + has shifted, so we never silently fall through to registering the local + sign hook on FSDP-managed parameters (which would double-sign their grads). + """ + for module in model.modules(): + if not isinstance(module, FSDPModule): + continue + try: + state = module._get_fsdp_state() + except AttributeError as exc: + raise RuntimeError(_FSDP_INTERNALS_ERROR) from exc + fsdp_param_group = getattr(state, "_fsdp_param_group", None) + if fsdp_param_group is None: + continue + try: + fsdp_params = fsdp_param_group.fsdp_params + except AttributeError as exc: + raise RuntimeError(_FSDP_INTERNALS_ERROR) from exc + for fsdp_param in fsdp_params: + try: + yield fsdp_param.sharded_param + except AttributeError as exc: + raise RuntimeError(_FSDP_INTERNALS_ERROR) from exc + + +def _get_fsdp_managed_param_ids(model: torch.nn.Module) -> set[int]: + return {id(p) for p in _iter_fsdp_managed_sharded_params(model)} + + +def _assert_no_fsdp_managed_local_hook(model: torch.nn.Module) -> None: + """Verify FSDP-managed parameters never received the local sign hook. + + Re-walking after registration is the explicit cross-check that catches + the failure mode where a torch upgrade rearranges the FSDP2 internals + we relied on but doesn't raise (e.g. an empty enumeration that lets + every parameter slip through to ``_register_local_sign_hooks``). + """ + for sharded in _iter_fsdp_managed_sharded_params(model): + if getattr(sharded, "_distsign_local_hook_registered", False): + raise RuntimeError( + "DistSignSGD installed the local sign hook on an FSDP-managed parameter, " + "which would double-sign gradients. " + _FSDP_INTERNALS_ERROR + ) + + +def _register_local_sign_hooks(model: torch.nn.Module, *, managed_param_ids: Optional[set[int]] = None) -> int: + """Sign plain-tensor gradients before they accumulate across microbatches.""" + count = 0 + managed_param_ids = managed_param_ids or set() + for param in model.parameters(): + if not param.requires_grad: + continue + if id(param) in managed_param_ids: + continue + if DTensor is not None and isinstance(param, DTensor): + # Non-FSDP DTensors (e.g. pure TP/PP shards) would get neither the + # local sign hook nor the FSDP custom reduce-scatter, so their + # gradients would skip the sign step entirely and pollute the + # voter total. Refuse rather than silently drift. + raise NotImplementedError( + "DistSignSGD encountered a DTensor parameter that is not FSDP-managed " + "(e.g. a pure TP/PP shard). This is not supported because such gradients " + "would bypass the sign step. Please file an issue if this configuration " + "is needed." + ) + if getattr(param, "_distsign_local_hook_registered", False): + continue + param.register_hook(_local_sign_hook) + param._distsign_local_hook_registered = True + count += 1 + return count + + +def _configure_fsdp_reduce_scatter(model: torch.nn.Module, *, sp_group: Optional[dist.ProcessGroup]) -> int: + count = 0 + for module in model.modules(): + if not isinstance(module, FSDPModule): + continue + if getattr(module, "_distsign_reduce_scatter_configured", False): + continue + module.set_custom_reduce_scatter(DistSignReduceScatter(sp_group=sp_group)) + module._distsign_reduce_scatter_configured = True + count += 1 + return count + + +def configure_distsignsgd(model: torch.nn.Module) -> None: + """Install DistSignSGD communication and local grad hooks on a parallelized model.""" + if getattr(model, "_distsignsgd_configured", False): + return + + ps = get_parallel_state() + # Direct attribute access (not getattr-with-default) so that an attribute + # rename on `ParallelState` surfaces as AttributeError instead of silently + # falling back to the default and skipping the guard. + if ps.dp_mode != "fsdp2": + raise ValueError("DistSignSGD requires data_parallel_mode='fsdp2'.") + if ps.dp_replicate_enabled: + raise NotImplementedError("DistSignSGD does not yet support HSDP / dp_replicate_size > 1.") + if ps.ep_enabled: + raise NotImplementedError( + "DistSignSGD does not yet support expert parallelism (EP). EP-managed grads " + "live on a different fsdp group, so a single active_voter_total cannot " + "normalize them correctly." + ) + if ps.cp_enabled and ps.cp_fsdp_mode != "none": + raise NotImplementedError( + "DistSignSGD does not support folding sequence-parallel exact-sum dims into FSDP; set cp_fsdp_mode='none'." + ) + + managed_param_ids = _get_fsdp_managed_param_ids(model) + _configure_fsdp_reduce_scatter(model, sp_group=ps.sp_grad_sync_group) + _register_local_sign_hooks(model, managed_param_ids=managed_param_ids) + _assert_no_fsdp_managed_local_hook(model) + model._distsignsgd_configured = True + + +class DistSignSGD(Optimizer): + """State-free SGD that expects gradients to be pre-signed before step().""" + + 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): + 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("DistSignSGD does not support sparse gradients.") + + if weight_decay: + p.add_(p, alpha=-lr * weight_decay) + + p.add_(grad, alpha=-lr) + + return loss diff --git a/src/xorl/optim/lr_scheduler.py b/src/xorl/optim/lr_scheduler.py index 62ec06ff..f4f6f991 100644 --- a/src/xorl/optim/lr_scheduler.py +++ b/src/xorl/optim/lr_scheduler.py @@ -50,6 +50,11 @@ def build_lr_scheduler( lr_min: float = 1e-7, lr_start: float = 0.0, ): + if lr <= 0: + raise ValueError(f"lr must be > 0, got {lr}.") + if not 0.0 <= lr_warmup_ratio <= 1.0: + raise ValueError(f"lr_warmup_ratio must be in [0, 1], got {lr_warmup_ratio}.") + # Handle MultiOptimizer by creating one scheduler per underlying optimizer if hasattr(optimizer, "_is_multi_optimizer") or isinstance(optimizer, dict): schedulers = {} @@ -81,6 +86,8 @@ def build_lr_scheduler( num_warmup_steps=lr_warmup_steps, num_training_steps=train_steps, init_lr=lr, + lr_decay_ratio=lr_decay_ratio, + min_lr=lr_min, lr_start=lr_start, ) @@ -125,23 +132,26 @@ def get_linear_schedule_with_warmup( num_training_steps: int, init_lr: float, last_epoch: int = -1, + lr_decay_ratio: float = 1.0, min_lr: float = 1e-7, lr_start: float = 0.0, ): """ - Creates a schedule with a learning rate that decreases linearly from the initial lr set in the optimizer to 0, - after a warmup period during which it increases linearly from 0 to the initial lr set in the optimizer. + Creates a schedule with a learning rate that decreases linearly from the initial lr set in the optimizer to + min_lr, after a warmup period during which it increases linearly from lr_start to the initial lr. """ + lr_decay_steps = int(num_training_steps * lr_decay_ratio) def _lr_lambda(current_step: int): if current_step < num_warmup_steps: return (lr_start + (init_lr - lr_start) * current_step / max(1, num_warmup_steps)) / init_lr min_lr_ratio = min_lr / init_lr - return max( - min_lr_ratio, - float(num_training_steps - current_step) / float(max(1, num_training_steps - num_warmup_steps)), - ) + if current_step >= lr_decay_steps: + return min_lr_ratio + + progress = float(current_step - num_warmup_steps) / float(max(1, lr_decay_steps - num_warmup_steps)) + return min_lr_ratio + (1.0 - min_lr_ratio) * (1.0 - progress) return LambdaLR(optimizer, _lr_lambda, last_epoch) @@ -163,19 +173,18 @@ def get_cosine_schedule_with_warmup( and the initial lr set in the optimizer. """ - def lr_lambda(current_step: int): - lr_decay_steps = int(num_training_steps * lr_decay_ratio) + lr_decay_steps = int(num_training_steps * lr_decay_ratio) + + def _lr_lambda(current_step: int): if current_step < num_warmup_steps: return (lr_start + (init_lr - lr_start) * current_step / max(1, num_warmup_steps)) / init_lr min_lr_ratio = min_lr / init_lr - if current_step > lr_decay_steps: + if current_step >= lr_decay_steps: return min_lr_ratio progress = float(current_step - num_warmup_steps) / float(max(1, lr_decay_steps - num_warmup_steps)) - assert 0 <= progress <= 1 factor = 0.5 * (1.0 + math.cos(math.pi * float(num_cycles) * 2.0 * progress)) - factor = factor * (1 - min_lr_ratio) + min_lr_ratio - return max(0, factor) + return min_lr_ratio + (1.0 - min_lr_ratio) * factor - return LambdaLR(optimizer, lr_lambda, last_epoch) + return LambdaLR(optimizer, _lr_lambda, last_epoch) diff --git a/src/xorl/optim/muon.py b/src/xorl/optim/muon.py index 1f84e1f2..e817d458 100644 --- a/src/xorl/optim/muon.py +++ b/src/xorl/optim/muon.py @@ -3,7 +3,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) + - FSDP2/EP DTensor support (shard-local or full-gradient Newton-Schulz) - 3D+ MoE expert tensor support (preserve leading dims as matrix batches) The core Muon algorithm is aligned with PyTorch's implementation: @@ -24,21 +24,102 @@ from collections import defaultdict from dataclasses import dataclass -from typing import Iterable, Optional, Tuple +from typing import TYPE_CHECKING, Iterable, Optional, Tuple import torch from torch.distributed._tensor import DTensor +from torch.distributed.tensor import Shard from torch.optim import Muon as TorchMuon from torch.optim._muon import _adjust_lr, _zeropower_via_newtonschulz from torch.optim.optimizer import Optimizer + +if TYPE_CHECKING: + from torch.distributed.device_mesh import DeviceMesh + from torch.distributed.tensor.placement_types import Placement + from ..utils import logging +from .cautious import apply_cautious_decay_ 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 +GROUPED_GRAM_NS_FP32_BYTE_LIMIT = 512 * 1024**2 + + +def _batched_zeropower_via_newtonschulz( + grad: torch.Tensor, + ns_coefficients: Tuple[float, float, float], + ns_steps: int, + eps: float, +) -> torch.Tensor: + """Batched Newton-Schulz on a stack of matrices ``[B, H, I]``. + + Equivalent to looping the upstream 2D ``_zeropower_via_newtonschulz`` over + the leading batch dim, but emits one bmm/baddbmm per NS step instead of one + matmul per (batch, step). On Qwen3.5-35B-A3B (40 MoE layers × 8 local + experts × 5 NS steps) this collapses ~14k kernel launches per optimizer + step into ~600, recovering the per-expert-Python-loop regression while + keeping the per-matrix math identical. + """ + if ns_steps >= 100: + raise ValueError("Number of steps must be less than 100 for computational efficiency") + if grad.ndim != 3: + raise ValueError(f"Batched NS expects a 3D tensor, got shape {tuple(grad.shape)}") + if len(ns_coefficients) != 3: + raise ValueError("Coefficients must be a tuple of exactly 3 values") + + a, b, c = ns_coefficients + + # Match upstream behavior: cast to bf16, optionally transpose so H <= I. + ortho = grad.bfloat16() + transposed = ortho.size(-2) > ortho.size(-1) + if transposed: + ortho = ortho.transpose(-2, -1).contiguous() + + # Per-matrix spectral-norm normalisation: divide each batch element by its + # own Frobenius norm (upper bound on spectral norm). Matches the upstream + # ``ortho_grad.div_(ortho_grad.norm().clamp(min=eps))``. + norms = ortho.flatten(start_dim=1).norm(dim=1).clamp(min=eps).reshape(-1, 1, 1) + ortho = ortho / norms + + for _ in range(ns_steps): + # gram_matrix[i] = ortho[i] @ ortho[i].T + gram_matrix = torch.bmm(ortho, ortho.transpose(-2, -1)) + # gram_update[i] = b * gram_matrix[i] + c * gram_matrix[i] @ gram_matrix[i] + gram_update = torch.baddbmm(gram_matrix, gram_matrix, gram_matrix, beta=b, alpha=c) + # ortho[i] = a * ortho[i] + gram_update[i] @ ortho[i] + ortho = torch.baddbmm(ortho, gram_update, ortho, beta=a) + + if transposed: + ortho = ortho.transpose(-2, -1) + return ortho + + +def _shard_full_to_local(full: torch.Tensor, mesh, placements) -> torch.Tensor: + """Slice a globally-replicated tensor down to its local shard for ``placements``. + + Mirrors DTensor's chunk-based ``Shard.split_tensor`` semantics: along each + sharded mesh dim, split with ``torch.chunk`` and select the local rank's + chunk. Trailing ranks may receive an empty chunk when the dim is not + evenly divisible (matching ``DTensor._local_tensor`` shape conventions). + Replicate placements pass through unchanged. + """ + out = full + for mesh_dim, placement in enumerate(placements): + if isinstance(placement, Shard): + dim = placement.dim + world = mesh.size(mesh_dim) + rank = mesh.get_local_rank(mesh_dim) + chunks = list(torch.chunk(out, world, dim=dim)) + if rank < len(chunks): + out = chunks[rank].contiguous() + else: + shape = list(out.shape) + shape[dim] = 0 + out = torch.empty(shape, dtype=out.dtype, device=out.device) + return out @dataclass @@ -47,6 +128,8 @@ class _MuonUpdatePlan: adjusted_lr: float orig_shape: Optional[torch.Size] pieces: list[Optional[torch.Tensor]] + placements: Optional[Tuple["Placement", ...]] = None + device_mesh: Optional["DeviceMesh"] = None @dataclass @@ -104,10 +187,23 @@ class Muon(TorchMuon): gram_newton_schulz_restart_iterations: Explicit Gram Newton-Schulz restart iteration indices. A value of ``2`` means restart after finishing the second iteration. + grouped_gram_ns_fp32_byte_limit: Maximum fp32 scratch bytes per grouped + Gram Newton-Schulz batch before splitting it into smaller chunks. 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 the parameter. + distributed_mode: How to handle Newton-Schulz on FSDP2/EP-sharded + DTensor params. ``"shard_local"`` (default) runs NS on each + rank's local shard — cheap but only an approximation of full + Muon. ``"full_gradient"`` all-gathers the post-momentum update + to the full matrix, runs NS on the full matrix on every rank + in the param's mesh (redundantly), and slices the + orthogonalized update back to the local shard. This recovers + the exact Muon update direction at the cost of a per-step + all-gather and replicated NS compute. Implements the dense + path of DeepSeek V4 §3.5.1 (without knapsack bucket assignment; + see followup work). Non-DTensor params are unaffected. """ def __init__( @@ -125,11 +221,14 @@ def __init__( grad_dtype: Optional[torch.dtype] = None, update_dtype: Optional[torch.dtype] = None, force_momentum_path: bool = False, - ns_algorithm: str = "standard_newton_schulz", + ns_algorithm: str = "gram_newton_schulz", ns_use_quack_kernels: bool = True, gram_newton_schulz_num_restarts: int = 1, gram_newton_schulz_restart_iterations: Optional[Iterable[int]] = None, + grouped_gram_ns_fp32_byte_limit: int = GROUPED_GRAM_NS_FP32_BYTE_LIMIT, adamw_state_dtype: Optional[torch.dtype] = None, + cautious: bool = False, + distributed_mode: str = "shard_local", ): if ns_algorithm not in {"standard_newton_schulz", "gram_newton_schulz"}: raise ValueError( @@ -140,12 +239,19 @@ def __init__( raise ValueError( f"gram_newton_schulz_num_restarts must be non-negative, got {gram_newton_schulz_num_restarts}" ) + if distributed_mode not in {"shard_local", "full_gradient"}: + raise ValueError( + f"Unsupported Muon distributed_mode: {distributed_mode!r}. Expected 'shard_local' or 'full_gradient'." + ) + if grouped_gram_ns_fp32_byte_limit <= 0: + raise ValueError(f"grouped_gram_ns_fp32_byte_limit must be positive, got {grouped_gram_ns_fp32_byte_limit}") 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._distributed_mode = distributed_mode self._logged_dtypes = False self._gram_ns_orthogonalizers = {} # Skip TorchMuon.__init__ (which enforces 2D-only) and call @@ -168,9 +274,11 @@ def __init__( if gram_newton_schulz_restart_iterations is not None else None ), + grouped_gram_ns_fp32_byte_limit=grouped_gram_ns_fp32_byte_limit, use_muon=True, adamw_betas=adamw_betas, adamw_eps=adamw_eps, + cautious=cautious, ) Optimizer.__init__(self, params, defaults) @@ -201,8 +309,15 @@ def _muon_step(self, group: dict) -> None: 5. Weight decay: param *= 1 - lr * wd 6. Update: param -= adjusted_lr * update - For FSDP2/EP DTensors, operates on the local shard directly - (shard-local Newton-Schulz) to avoid DTensor reshape/matmul issues. + For FSDP2/EP DTensors, ``distributed_mode`` selects between two paths: + - ``"shard_local"`` (default): NS is applied to each rank's local + shard. Cheap; an approximation of full Muon. + - ``"full_gradient"``: post-momentum update is all-gathered to the + full matrix, NS is applied on the full matrix on every rank in + the param's mesh, and the orthogonalized update is sliced back + to the local shard before being written to ``p._local_tensor``. + Recovers exact Muon at the cost of a per-step all-gather and + replicated NS compute. """ lr = group["lr"] momentum = group["momentum"] @@ -212,6 +327,7 @@ def _muon_step(self, group: dict) -> None: eps = group["eps"] weight_decay = group["weight_decay"] adjust_lr_fn = group["adjust_lr_fn"] + cautious = group.get("cautious", False) 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) @@ -224,6 +340,11 @@ def _muon_step(self, group: dict) -> None: # Extract local tensors from DTensors (FSDP2/EP sharded params). grad = p.grad is_dtensor = isinstance(grad, DTensor) + use_full_gradient = ( + self._distributed_mode == "full_gradient" + and is_dtensor + and any(isinstance(pl, Shard) for pl in grad.placements) + ) if is_dtensor: grad_local = grad._local_tensor p_local = p._local_tensor @@ -236,9 +357,12 @@ def _muon_step(self, group: dict) -> None: # Handle 3D+ tensors as batches of matrices: [..., hidden, intermediate]. # For fused gate_up_proj [E, H, 2I], split into two [..., H, I] halves. + # In shard_local mode this runs on the local shard; in full_gradient + # mode we defer the reshape until after the post-momentum all-gather + # so the reshape and LR adjustment see the full matrix shape. orig_shape = None fused_split = None - if grad_local.ndim >= 3: + if not use_full_gradient and grad_local.ndim >= 3: orig_shape = grad_local.shape fused_gate_up_ids = group.get("_fused_gate_up_ids", set()) if id(p) in fused_gate_up_ids: @@ -285,10 +409,30 @@ def _muon_step(self, group: dict) -> None: ) self._logged_dtypes = True + # Full-gradient mode: all-gather the post-momentum update to the + # full matrix shape on every rank in the param's mesh, then run + # NS on the full matrix. Momentum/Nesterov are linear in the + # gradient and commute with sharding, so doing them on the local + # shard before gather is mathematically identical to doing them + # post-gather but uses only local-shard buffer memory. + plan_placements = None + plan_mesh = None + if use_full_gradient: + plan_placements = grad.placements + plan_mesh = grad.device_mesh + update_dtensor = DTensor.from_local(update, plan_mesh, plan_placements, run_check=False) + update = update_dtensor.full_tensor() + if update.ndim >= 3: + orig_shape = update.shape + fused_gate_up_ids = group.get("_fused_gate_up_ids", set()) + if id(p) in fused_gate_up_ids: + fused_split = update.shape[-1] // 2 + update = update.reshape(-1, *update.shape[-2:]) + adjusted_lr = _adjust_lr( lr, adjust_lr_fn, - grad_local.shape[-2:] if grad_local.ndim > 2 else grad_local.shape, + update.shape[-2:] if update.ndim > 2 else update.shape, ) pieces = [update[..., :fused_split], update[..., fused_split:]] if fused_split is not None else [update] plan = _MuonUpdatePlan( @@ -296,6 +440,8 @@ def _muon_step(self, group: dict) -> None: adjusted_lr=adjusted_lr, orig_shape=orig_shape, pieces=[None] * len(pieces), + placements=plan_placements, + device_mesh=plan_mesh, ) update_plans.append(plan) @@ -320,7 +466,11 @@ def _muon_step(self, group: dict) -> None: if uses_grouped_gram_ns and grouped_updates: orthogonalizer = self._get_gram_ns_orthogonalizer(group) - self._orthogonalize_grouped_gram_ns_updates(grouped_updates, orthogonalizer) + self._orthogonalize_grouped_gram_ns_updates( + grouped_updates, + orthogonalizer, + group["grouped_gram_ns_fp32_byte_limit"], + ) for plan in update_plans: update_pieces = plan.pieces @@ -335,11 +485,23 @@ def _muon_step(self, group: dict) -> None: if plan.orig_shape is not None: update = update.reshape(plan.orig_shape) + # Slice the orthogonalized full-tensor update back to the local + # shard for full_gradient mode. This is purely local (no comm). + if plan.placements is not None: + update = _shard_full_to_local(update, plan.device_mesh, plan.placements) + # Cast back to param dtype update = update.to(plan.param.dtype) - # Decoupled weight decay - plan.param.mul_(1 - lr * weight_decay) + # Decoupled weight decay (cautious mask uses the post-NS update, + # which is the actual u_t direction for Muon). + apply_cautious_decay_( + plan.param, + update_sign_proxy=update, + lr=lr, + weight_decay=weight_decay, + cautious=cautious, + ) # Parameter update plan.param.add_(update, alpha=-plan.adjusted_lr) @@ -358,10 +520,9 @@ def _orthogonalize_update( 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) + return _batched_zeropower_via_newtonschulz(flat_update, ns_coefficients, ns_steps, eps).reshape( + original_shape + ) if group["ns_algorithm"] == "gram_newton_schulz": return self._get_gram_ns_orthogonalizer(group).orthogonalize(update) raise ValueError( @@ -373,9 +534,10 @@ def _orthogonalize_grouped_gram_ns_updates( self, grouped_updates: dict[tuple[tuple[int, int], torch.dtype, torch.device], list[_GroupedOrthogonalizationEntry]], orthogonalizer: GramNewtonSchulzOrthogonalizer, + fp32_byte_limit: int, ) -> None: for batch_entries in grouped_updates.values(): - for chunk in self._iter_grouped_gram_ns_chunks(batch_entries): + for chunk in self._iter_grouped_gram_ns_chunks(batch_entries, fp32_byte_limit): 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) @@ -398,9 +560,10 @@ def _orthogonalize_grouped_gram_ns_updates( def _iter_grouped_gram_ns_chunks( self, batch_entries: list[_GroupedOrthogonalizationEntry], + fp32_byte_limit: int, ): 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)) + max_matrix_batch = max(1, fp32_byte_limit // (per_matrix_numel * 4)) chunk: list[_GroupedOrthogonalizationEntry] = [] chunk_matrix_batch = 0 @@ -457,6 +620,7 @@ def _adamw_step(self, group: dict) -> None: beta1, beta2 = group["adamw_betas"] eps = group["adamw_eps"] weight_decay = group["weight_decay"] + cautious = group.get("cautious", False) for p in group["params"]: if p.grad is None: @@ -478,14 +642,20 @@ def _adamw_step(self, group: dict) -> None: exp_avg = state["exp_avg"] exp_avg_sq = state["exp_avg_sq"] - # Decoupled weight decay - if weight_decay > 0: - p.data.mul_(1.0 - lr * weight_decay) - # Update biased first and second moment estimates exp_avg.mul_(beta1).add_(grad, alpha=1.0 - beta1) exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1.0 - beta2) + # Decoupled weight decay (cautious mask uses sign(exp_avg) which + # matches sign(u_t) since the Adam denominator is positive). + apply_cautious_decay_( + p.data, + update_sign_proxy=exp_avg, + lr=lr, + weight_decay=weight_decay, + cautious=cautious, + ) + # Bias correction bias_correction1 = 1.0 - beta1**step bias_correction2 = 1.0 - beta2**step diff --git a/src/xorl/optim/optimizer.py b/src/xorl/optim/optimizer.py index 84ddf988..f09794e2 100644 --- a/src/xorl/optim/optimizer.py +++ b/src/xorl/optim/optimizer.py @@ -12,8 +12,9 @@ from ..distributed.parallel_state import get_parallel_state from ..utils import logging from .anyprecision_adamw import AnyPrecisionAdamW +from .distsignsgd import DistSignSGD, configure_distsignsgd from .multi_optimizer import MultiOptimizer -from .muon import Muon +from .muon import GROUPED_GRAM_NS_FP32_BYTE_LIMIT, Muon from .signsgd import SignSGD @@ -114,11 +115,42 @@ def _get_optimizer_cls_and_kwargs( fused: bool = False, optimizer_dtype: str = "bf16", optimizer_kwargs: Optional[Dict[str, Any]] = None, + cautious_weight_decay: bool = False, ) -> Tuple[type, Dict[str, Any]]: """Return (optimizer_class, constructor_kwargs) without instantiating.""" kwargs = optimizer_kwargs or {} if optimizer_type == "adamw": + if cautious_weight_decay: + # torch.optim.AdamW has no cautious-decay hook; route to our + # AnyPrecisionAdamW (fp32 state -> mathematically equivalent to + # torch AdamW) which supports the mask. + logger.info_rank0( + "cautious_weight_decay=True with optimizer=adamw: routing to AnyPrecisionAdamW (fp32 state)." + ) + # Reject torch.optim.AdamW-only kwargs up-front: forwarding them + # to AnyPrecisionAdamW yields a confusing TypeError naming a class + # the user did not request. + _ANYPRECISION_ACCEPTED = {"use_kahan_summation"} + unsupported = [k for k in kwargs if k not in _ANYPRECISION_ACCEPTED] + if unsupported: + raise ValueError( + f"cautious_weight_decay=True with optimizer='adamw' routes to AnyPrecisionAdamW, " + f"which does not accept these optimizer_kwargs: {unsupported}. " + f"Either drop them, set optimizer='anyprecision_adamw' explicitly, or disable cautious." + ) + ctor_kwargs = dict( + lr=lr, + betas=betas, + eps=eps, + weight_decay=weight_decay, + momentum_dtype=torch.float32, + variance_dtype=torch.float32, + compensation_buffer_dtype=torch.float32, + cautious=True, + **kwargs, + ) + return AnyPrecisionAdamW, ctor_kwargs foreach = not fused ctor_kwargs = dict( lr=lr, @@ -140,17 +172,27 @@ def _get_optimizer_cls_and_kwargs( momentum_dtype=state_dtype, variance_dtype=state_dtype, compensation_buffer_dtype=state_dtype, + cautious=cautious_weight_decay, **kwargs, ) return AnyPrecisionAdamW, ctor_kwargs elif optimizer_type == "sgd": + if cautious_weight_decay: + raise ValueError( + "cautious_weight_decay is not supported with optimizer='sgd' " + "(torch.optim.SGD has no per-coordinate decay hook). " + "Use 'anyprecision_adamw', 'signsgd', or 'muon' instead." + ) sgd_defaults = {"momentum": 0.0, "nesterov": False} 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) + ctor_kwargs = dict(lr=lr, weight_decay=weight_decay, cautious=cautious_weight_decay, **kwargs) return SignSGD, ctor_kwargs + elif optimizer_type == "distsignsgd": + ctor_kwargs = dict(lr=lr, weight_decay=weight_decay, **kwargs) + return DistSignSGD, ctor_kwargs elif optimizer_type == "muon": adamw_state_dtype = _ANYPRECISION_STATE_DTYPES.get(optimizer_dtype) momentum_dtype = kwargs.get("muon_momentum_dtype") @@ -163,10 +205,14 @@ 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_algorithm=kwargs.get("muon_ns_algorithm", "gram_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"), + grouped_gram_ns_fp32_byte_limit=kwargs.get( + "muon_grouped_gram_ns_fp32_byte_limit", + GROUPED_GRAM_NS_FP32_BYTE_LIMIT, + ), adamw_betas=betas, adamw_eps=eps, momentum_dtype=_normalize_optional_dtype(momentum_dtype, field_name="muon_momentum_dtype"), @@ -174,11 +220,13 @@ def _get_optimizer_cls_and_kwargs( 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, + cautious=cautious_weight_decay, + distributed_mode=kwargs.get("muon_distributed_mode", "shard_local"), ) return Muon, ctor_kwargs else: raise ValueError( - f"Unsupported optimizer type: '{optimizer_type}'. Supported: adamw, anyprecision_adamw, sgd, signsgd, muon." + f"Unsupported optimizer type: '{optimizer_type}'. Supported: adamw, anyprecision_adamw, sgd, signsgd, distsignsgd, muon." ) @@ -193,6 +241,7 @@ def _create_optimizer( fused: bool = False, optimizer_dtype: str = "bf16", optimizer_kwargs: Optional[Dict[str, Any]] = None, + cautious_weight_decay: bool = False, ) -> Optimizer: """ Single factory for all optimizer types. @@ -201,6 +250,7 @@ def _create_optimizer( Optimizer-specific args are passed via optimizer_kwargs: - sgd: {"momentum": 0.9, "nesterov": True} - signsgd: no optimizer-specific kwargs + - distsignsgd: no optimizer-specific kwargs - muon: {"muon_lr": 0.02, "muon_momentum": 0.95, ...} - adamw/anyprecision_adamw: any extra kwargs forwarded to constructor """ @@ -213,6 +263,7 @@ def _create_optimizer( fused=fused, optimizer_dtype=optimizer_dtype, optimizer_kwargs=optimizer_kwargs, + cautious_weight_decay=cautious_weight_decay, ) return cls(param_groups, **ctor_kwargs) @@ -332,6 +383,7 @@ def build_optimizer( no_decay_modules: Optional[List[str]] = None, no_decay_params: Optional[List[str]] = None, optimizer_kwargs: Optional[Dict[str, Any]] = None, + cautious_weight_decay: bool = False, ) -> "torch.optim.Optimizer": """ Build an optimizer for the given model. @@ -343,7 +395,7 @@ 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", "signsgd", "muon". + optimizer_type: One of "adamw", "anyprecision_adamw", "sgd", "signsgd", "distsignsgd", "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. @@ -351,17 +403,28 @@ def build_optimizer( optimizer_kwargs: Optimizer-specific keyword arguments passed to the constructor. - sgd: {"momentum": 0.9, "nesterov": True} - signsgd: no optimizer-specific kwargs + - distsignsgd: 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_ns_algorithm": "standard_newton_schulz", + "muon_ns_algorithm": "gram_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} + "muon_grad_dtype": None, "muon_update_dtype": None, "muon_force_momentum_path": False, + "muon_distributed_mode": "shard_local"} - adamw/anyprecision_adamw: any extra kwargs forwarded to constructor + cautious_weight_decay: If True, apply Cautious Weight Decay (CWD) per + Chen et al. (arXiv:2510.12402): mask the decoupled decay term by + ``I(u_t * x_t >= 0)``. Supported for adamw, anyprecision_adamw, + signsgd, and muon. With ``optimizer_type='adamw'`` this routes to + AnyPrecisionAdamW with fp32 state (no fused kernel). """ + ps = get_parallel_state() + if optimizer_type == "distsignsgd" and ps.dp_mode != "fsdp2": + raise ValueError("DistSignSGD requires data_parallel_mode='fsdp2'.") + # 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( + optimizer = build_ep_fsdp2_optimizer( model, lr, betas, @@ -374,7 +437,11 @@ def build_optimizer( no_decay_modules, no_decay_params, optimizer_kwargs=optimizer_kwargs, + cautious_weight_decay=cautious_weight_decay, ) + if optimizer_type == "distsignsgd": + configure_distsignsgd(model) + return optimizer kwargs = optimizer_kwargs or {} @@ -415,7 +482,7 @@ def build_optimizer( logger.info_rank0(f"Parameters without weight decay: {no_decay_parameter_names}") param_groups.append({"params": no_decay_parameters, "weight_decay": 0.0}) - return _create_optimizer( + optimizer = _create_optimizer( optimizer_type, param_groups, lr=lr, @@ -425,7 +492,11 @@ def build_optimizer( fused=fused, optimizer_dtype=optimizer_dtype, optimizer_kwargs=optimizer_kwargs, + cautious_weight_decay=cautious_weight_decay, ) + if optimizer_type == "distsignsgd": + configure_distsignsgd(model) + return optimizer def build_ep_fsdp2_optimizer( @@ -441,6 +512,7 @@ def build_ep_fsdp2_optimizer( no_decay_modules: Optional[List[str]] = None, no_decay_params: Optional[List[str]] = None, optimizer_kwargs: Optional[Dict[str, Any]] = None, + cautious_weight_decay: bool = False, ): """ Build a MultiOptimizer instance when model is parallelized with EP+FSDP2 @@ -524,6 +596,7 @@ def _build(groups: Sequence[Dict[str, Any]]) -> Optimizer: fused=fused, optimizer_dtype=optimizer_dtype, optimizer_kwargs=optimizer_kwargs, + cautious_weight_decay=cautious_weight_decay, ) optimizer_dict: Dict[str, Optimizer] = {} diff --git a/src/xorl/optim/signsgd.py b/src/xorl/optim/signsgd.py index 9de7eab7..7eaef54d 100644 --- a/src/xorl/optim/signsgd.py +++ b/src/xorl/optim/signsgd.py @@ -1,16 +1,24 @@ import torch from torch.optim.optimizer import Optimizer +from .cautious import apply_cautious_decay_ + # 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.""" + """Sign-based SGD optimizer with no optimizer-state tensors. + + When ``cautious=True``, decoupled weight decay is masked by + ``I(sign(grad) * param >= 0)`` per Chen et al. "Cautious Weight Decay" + (arXiv:2510.12402). Equivalently, the mask is ``I(grad * param >= 0)``. + """ def __init__( self, params, lr: float = 1e-3, weight_decay: float = 0.0, + cautious: bool = False, ): if lr < 0.0: raise ValueError(f"Invalid learning rate: {lr}") @@ -20,6 +28,7 @@ def __init__( defaults = { "lr": lr, "weight_decay": weight_decay, + "cautious": cautious, } super().__init__(params, defaults) @@ -34,6 +43,7 @@ def step(self, closure=None): for group in self.param_groups: lr = group["lr"] weight_decay = group["weight_decay"] + cautious = group.get("cautious", False) for p in group["params"]: grad = p.grad @@ -42,8 +52,13 @@ def step(self, closure=None): if grad.is_sparse: raise RuntimeError("SignSGD does not support sparse gradients.") - if weight_decay: - p.add_(p, alpha=-lr * weight_decay) + apply_cautious_decay_( + p, + update_sign_proxy=grad, + lr=lr, + weight_decay=weight_decay, + cautious=cautious, + ) p.add_(torch.sign(grad), alpha=-lr) diff --git a/src/xorl/optim/stochastic_round.py b/src/xorl/optim/stochastic_round.py new file mode 100644 index 00000000..777379dc --- /dev/null +++ b/src/xorl/optim/stochastic_round.py @@ -0,0 +1,43 @@ +"""Stochastic rounding to BF16. + +Stochastic rounding produces an unbiased estimator of an FP32 input: +``E[stochastic_round_to_bf16(x).to(fp32)] == x``. Round-to-nearest is biased +when accumulating many small values; stochastic rounding distributes the +rounding decision in proportion to where ``x`` falls between its two BF16 +neighbors, which is what makes BF16-transit reductions safe in expectation. + +The implementation manipulates the FP32 bit pattern: BF16 is the FP32 bit +pattern with the low 16 mantissa bits truncated. Adding a uniform random +integer in ``[0, 2**16)`` to those low bits before truncation gives, for an +FP32 value at fractional position ``f`` between BF16 neighbors, probability +``f`` of rounding up and ``1-f`` of rounding down. This is the standard +implementation used in TransformerEngine and torchao. +""" + +from typing import Optional + +import torch + + +def stochastic_round_to_bf16( + x: torch.Tensor, + *, + generator: Optional[torch.Generator] = None, +) -> torch.Tensor: + """Stochastically round an FP32 tensor to BF16. + + Args: + x: FP32 tensor. + generator: Optional ``torch.Generator`` for reproducibility. + + Returns: + BF16 tensor with the same shape and device as ``x``. + """ + if x.dtype != torch.float32: + raise ValueError(f"stochastic_round_to_bf16 requires fp32 input, got {x.dtype}") + + x_int = x.contiguous().view(torch.int32) + noise = torch.empty_like(x_int) + noise.random_(0, 1 << 16, generator=generator) + rounded_int = (x_int + noise) & ~0xFFFF + return rounded_int.view(torch.float32).to(torch.bfloat16) diff --git a/src/xorl/qlora/modules/linear.py b/src/xorl/qlora/modules/linear.py index 00daed2f..849dddd7 100644 --- a/src/xorl/qlora/modules/linear.py +++ b/src/xorl/qlora/modules/linear.py @@ -79,6 +79,8 @@ def __init__( # LoRA parameters (trainable, fp32) self.r = r self.lora_alpha = lora_alpha + self.active_r = r + self.active_lora_alpha = lora_alpha self.scaling = lora_alpha / r self.lora_A = nn.Parameter(torch.empty(r, in_features, dtype=torch.float32, device=device)) self.lora_B = nn.Parameter(torch.empty(out_features, r, dtype=torch.float32, device=device)) @@ -104,6 +106,16 @@ def reset_lora_parameters(self) -> None: nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5)) nn.init.zeros_(self.lora_B) + def set_runtime_lora_config(self, lora_rank: int, lora_alpha: int) -> None: + """Update the active LoRA slice used during forward/merge/export.""" + if lora_rank <= 0 or lora_rank > self.r: + raise ValueError(f"Active LoRA rank must be in [1, {self.r}], got {lora_rank}") + self.active_r = lora_rank + self.active_lora_alpha = lora_alpha + + def _active_scaling(self) -> float: + return self.active_lora_alpha / self.active_r + # ------------------------------------------------------------------ # Construction: dispatch to subclass based on quant_format # ------------------------------------------------------------------ @@ -364,7 +376,9 @@ def forward(self, x: Tensor) -> Tensor: result = F.linear(x, w, self.bias) x_lora = x.to(self.lora_A.dtype) - lora_out = F.linear(F.linear(x_lora, self.lora_A), self.lora_B) * self.scaling + lora_A = self.lora_A[: self.active_r] + lora_B = self.lora_B[:, : self.active_r] + lora_out = F.linear(F.linear(x_lora, lora_A), lora_B) * self._active_scaling() return result + lora_out.to(result.dtype) @@ -379,9 +393,9 @@ def _to_regular_tensor(t: Tensor) -> Tensor: return t def get_delta_weight(self) -> Tensor: - lora_A = self._to_regular_tensor(self.lora_A) - lora_B = self._to_regular_tensor(self.lora_B) - return (lora_B @ lora_A) * self.scaling + lora_A = self._to_regular_tensor(self.lora_A[: self.active_r]) + lora_B = self._to_regular_tensor(self.lora_B[:, : self.active_r]) + return (lora_B @ lora_A) * self._active_scaling() # ------------------------------------------------------------------ # Checkpoint utilities diff --git a/src/xorl/qlora/modules/moe_experts.py b/src/xorl/qlora/modules/moe_experts.py index 2edf0e14..7f651f02 100644 --- a/src/xorl/qlora/modules/moe_experts.py +++ b/src/xorl/qlora/modules/moe_experts.py @@ -83,6 +83,8 @@ def __init__( self.hidden_dim = hidden_size # alias for compatibility with MoEExpertsLoRA self.r = r self.lora_alpha = lora_alpha + self.active_r = r + self.active_lora_alpha = lora_alpha self.scaling = compute_lora_scaling(lora_alpha, r, use_rslora) self.quant_format = quant_format self.quant_group_size = quant_group_size @@ -194,14 +196,28 @@ def reset_lora_parameters(self): nn.init.kaiming_uniform_(lora_A.data[i], a=math.sqrt(5)) nn.init.zeros_(lora_B.data) + def set_runtime_lora_config(self, lora_rank: int, lora_alpha: int) -> None: + """Update the active LoRA slice used during forward/merge/export.""" + if lora_rank <= 0 or lora_rank > self.r: + raise ValueError(f"Active LoRA rank must be in [1, {self.r}], got {lora_rank}") + self.active_r = lora_rank + self.active_lora_alpha = lora_alpha + + def _active_scaling(self) -> float: + return compute_lora_scaling(self.active_lora_alpha, self.active_r, self.use_rslora) + + def _active_lora_views(self, proj_name: str) -> tuple[Tensor, Tensor]: + lora_A = getattr(self, f"{proj_name}_lora_A")[..., : self.active_r].contiguous() + lora_B = getattr(self, f"{proj_name}_lora_B")[:, : self.active_r, ...].contiguous() + return lora_A, lora_B + def _compute_proj_delta(self, proj_name: str) -> torch.Tensor: """Compute LoRA delta for one projection. Returns [E, K, N] in GKN format.""" - lora_A = getattr(self, f"{proj_name}_lora_A") # [1 or E, in, r] - lora_B = getattr(self, f"{proj_name}_lora_B") # [E or 1, r, out] + lora_A, lora_B = self._active_lora_views(proj_name) E = max(lora_A.shape[0], lora_B.shape[0]) A = lora_A.expand(E, -1, -1) # [E, in, r] B = lora_B.expand(E, -1, -1) # [E, r, out] - return torch.bmm(A, B) * self.scaling # [E, in, out] = [E, K, N] + return torch.bmm(A, B) * self._active_scaling() # [E, in, out] = [E, K, N] # ------------------------------------------------------------------ # Abstract methods (subclasses must implement) @@ -475,13 +491,13 @@ def forward( gate_proj=self.gate_proj.to(compute_dtype), up_proj=self.up_proj.to(compute_dtype), down_proj=self.down_proj.to(compute_dtype), - gate_proj_lora_A=self.gate_proj_lora_A, - gate_proj_lora_B=self.gate_proj_lora_B, - up_proj_lora_A=self.up_proj_lora_A, - up_proj_lora_B=self.up_proj_lora_B, - down_proj_lora_A=self.down_proj_lora_A, - down_proj_lora_B=self.down_proj_lora_B, - scaling=self.scaling, + gate_proj_lora_A=self._active_lora_views("gate_proj")[0], + gate_proj_lora_B=self._active_lora_views("gate_proj")[1], + up_proj_lora_A=self._active_lora_views("up_proj")[0], + up_proj_lora_B=self._active_lora_views("up_proj")[1], + down_proj_lora_A=self._active_lora_views("down_proj")[0], + down_proj_lora_B=self._active_lora_views("down_proj")[1], + scaling=self._active_scaling(), ) @torch.compiler.disable @@ -519,21 +535,28 @@ def _ep_forward( # Step 2: Expert computation with dequantized base + LoRA compute_dtype = permute_tokens.dtype + gate_proj_lora_A, gate_proj_lora_B = self._active_lora_views("gate_proj") + up_proj_lora_A, up_proj_lora_B = self._active_lora_views("up_proj") + down_proj_lora_A, down_proj_lora_B = self._active_lora_views("down_proj") expert_output = compute_fn( 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.scaling, + gate_proj_lora_A, + gate_proj_lora_B, + up_proj_lora_A, + up_proj_lora_B, + down_proj_lora_A, + down_proj_lora_B, + self._active_scaling(), ) + expert_scores = getattr(ctx, "expert_scores", getattr(ctx, "permuted_scores", None)) + if expert_scores is not None: + expert_output = expert_output * expert_scores.unsqueeze(1).to(expert_output.dtype) + # Step 3: Combine expert outputs back to original ranks combine_kwargs = self._build_combine_kwargs(expert_output, ctx, dispatch_kwargs, parallel_state) return combine_fn(**combine_kwargs) @@ -584,24 +607,28 @@ def _eager_lora_forward(self, hidden_states: Tensor, expert_idx: int) -> Tensor: return torch.zeros_like(hidden_states[:, : self.hidden_size]) compute_dtype = hidden_states.dtype + active_scaling = self._active_scaling() + gate_A, gate_B = self._active_lora_views("gate_proj") + up_A, up_B = self._active_lora_views("up_proj") + down_A, down_B = self._active_lora_views("down_proj") # gate_proj: x @ W (no transpose with G,K,N) gate_w = self.dequantize_expert("gate", local_idx, self.hidden_size, self.intermediate_size) gate_out = torch.matmul(hidden_states, gate_w.to(compute_dtype)) # gate LoRA: (x @ A) @ B * scaling -- hybrid shared via min() - A = self.gate_proj_lora_A[min(local_idx, self.gate_proj_lora_A.shape[0] - 1)].to(compute_dtype) - B = self.gate_proj_lora_B[local_idx].to(compute_dtype) - gate_out = gate_out + torch.matmul(torch.matmul(hidden_states, A), B) * self.scaling + A = gate_A[min(local_idx, gate_A.shape[0] - 1)].to(compute_dtype) + B = gate_B[local_idx].to(compute_dtype) + gate_out = gate_out + torch.matmul(torch.matmul(hidden_states, A), B) * active_scaling # up_proj: x @ W up_w = self.dequantize_expert("up", local_idx, self.hidden_size, self.intermediate_size) up_out = torch.matmul(hidden_states, up_w.to(compute_dtype)) # up LoRA - A = self.up_proj_lora_A[min(local_idx, self.up_proj_lora_A.shape[0] - 1)].to(compute_dtype) - B = self.up_proj_lora_B[local_idx].to(compute_dtype) - up_out = up_out + torch.matmul(torch.matmul(hidden_states, A), B) * self.scaling + A = up_A[min(local_idx, up_A.shape[0] - 1)].to(compute_dtype) + B = up_B[local_idx].to(compute_dtype) + up_out = up_out + torch.matmul(torch.matmul(hidden_states, A), B) * active_scaling # Activation out = self.act_fn(gate_out) * up_out @@ -611,9 +638,9 @@ def _eager_lora_forward(self, hidden_states: Tensor, expert_idx: int) -> Tensor: down_out = torch.matmul(out, down_w.to(compute_dtype)) # down LoRA - A = self.down_proj_lora_A[local_idx].to(compute_dtype) - B = self.down_proj_lora_B[min(local_idx, self.down_proj_lora_B.shape[0] - 1)].to(compute_dtype) - down_out = down_out + torch.matmul(torch.matmul(out, A), B) * self.scaling + A = down_A[local_idx].to(compute_dtype) + B = down_B[min(local_idx, down_B.shape[0] - 1)].to(compute_dtype) + down_out = down_out + torch.matmul(torch.matmul(out, A), B) * active_scaling return down_out @@ -624,7 +651,7 @@ def extra_repr(self) -> str: f"expert_offset={self.expert_offset}, " f"intermediate_size={self.intermediate_size}, " f"hidden_size={self.hidden_size}, " - f"r={self.r}, quant_format={self.quant_format}, " + f"r={self.active_r}, max_r={self.r}, quant_format={self.quant_format}, " f"moe_implementation={self.moe_implementation}, " f"hybrid_shared={self.hybrid_shared}, " f"ep_dispatch={self.ep_dispatch}" diff --git a/src/xorl/qlora/utils.py b/src/xorl/qlora/utils.py index 613920f5..e8271697 100644 --- a/src/xorl/qlora/utils.py +++ b/src/xorl/qlora/utils.py @@ -377,7 +377,7 @@ def maybe_quantize_qlora(model: nn.Module) -> int: def maybe_load_and_quantize_moe_qlora( model: nn.Module, weights_path: str, - load_mode: str = "broadcast", + load_mode: str = "grouped", ) -> int: """Load bf16 MoE expert weights from checkpoint and quantize. @@ -387,7 +387,8 @@ def maybe_load_and_quantize_moe_qlora( Args: model: Model with QLoRAMoeExperts modules. weights_path: Path to HF model directory with bf16 weights. - load_mode: "broadcast" or "all_ranks". + load_mode: "grouped" or "all_ranks". Deferred QLoRA loading uses + rank-0 fanout for "grouped" as the compatibility fallback. Returns: Number of MoE modules loaded. @@ -415,7 +416,18 @@ def maybe_load_and_quantize_moe_qlora( "MoE expert loading requires model.safetensors.index.json." ) - shard_cache: dict = {} + use_broadcast = ( + load_mode == "grouped" and dist.is_available() and dist.is_initialized() and dist.get_world_size() > 1 + ) + if use_broadcast: + needed_shards = sorted(set(weight_map.values())) + logger.info( + f"Broadcasting {len(needed_shards)} shard files from rank 0 " + f"(rank={dist.get_rank()}, world={dist.get_world_size()})" + ) + shard_cache = _broadcast_shard_cache(needed_shards, weights_path) + else: + shard_cache = {} moe_count = 0 for module in model.modules(): @@ -540,7 +552,7 @@ def detect_prequantized_nvfp4(weights_path: str) -> bool: True if the checkpoint is pre-quantized NVFP4. """ - from xorl.models.checkpoint_handlers.buffers import detect_prequantized_checkpoint + from xorl.models.checkpoint_handlers.buffers import detect_prequantized_checkpoint # noqa: PLC0415 return detect_prequantized_checkpoint(weights_path) @@ -556,7 +568,7 @@ def detect_prequantized_block_fp8(weights_path: str) -> bool: Returns: True if the checkpoint is pre-quantized block FP8. """ - from xorl.models.checkpoint_handlers.buffers import detect_prequantized_block_fp8_checkpoint + from xorl.models.checkpoint_handlers.buffers import detect_prequantized_block_fp8_checkpoint # noqa: PLC0415 return detect_prequantized_block_fp8_checkpoint(weights_path) @@ -638,7 +650,7 @@ def _read_shard(shard_file): return shard_cache -def maybe_load_prequantized_qlora(model: nn.Module, weights_path: str, load_mode: str = "broadcast") -> int: +def maybe_load_prequantized_qlora(model: nn.Module, weights_path: str, load_mode: str = "grouped") -> int: """Load pre-quantized weights into QLoRALinear modules from checkpoint. For pre-quantized NVFP4 checkpoints (modelopt format), this replaces the @@ -647,8 +659,11 @@ def maybe_load_prequantized_qlora(model: nn.Module, weights_path: str, load_mode Also handles QLoRAMoeExperts (auto-detected internally via weight_map probing). - When distributed is initialized, uses rank 0 broadcast to avoid redundant - disk reads across ranks (~6-8× faster for large models). + When distributed is initialized, grouped mode uses rank-0 fanout to avoid + redundant disk reads across ranks (~6-8× faster for large models). The main + checkpoint handler performs true grouped fanout when QLoRA weights are + loaded inline; this deferred path keeps rank-0 fanout as the compatibility + fallback. Args: model: Model with QLoRALinear/QLoRAMoeExperts modules @@ -676,10 +691,10 @@ def maybe_load_prequantized_qlora(model: nn.Module, weights_path: str, load_mode ) # Loading mode: - # "broadcast" (default): rank 0 reads, broadcasts via NCCL. Best for shared/NFS filesystems. + # "grouped" (default): use rank-0 fanout in this deferred QLoRA path. # "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 == "grouped" and dist.is_available() and dist.is_initialized() and dist.get_world_size() > 1 ) if use_broadcast: needed_shards = sorted(set(weight_map.values())) diff --git a/src/xorl/server/api_server/__init__.py b/src/xorl/server/api_server/__init__.py index b63dd862..5c8f75e4 100644 --- a/src/xorl/server/api_server/__init__.py +++ b/src/xorl/server/api_server/__init__.py @@ -17,7 +17,9 @@ HealthCheckResponse, LoadWeightsRequest, LoadWeightsResponse, + LoRAConfigRequest, LossFnOutput, + OptimizerConfigRequest, OptimStepRequest, OptimStepResponse, SaveWeightsForSamplerRequest, @@ -34,6 +36,7 @@ "OrchestratorClient", "APIServer", # Request/Response models + "AdamParams", "TensorData", "Datum", "DatumInput", @@ -42,7 +45,8 @@ "ForwardBackwardRequest", "ForwardBackwardResponse", "LossFnOutput", - "AdamParams", + "LoRAConfigRequest", + "OptimizerConfigRequest", "OptimStepRequest", "OptimStepResponse", "SaveWeightsRequest", diff --git a/src/xorl/server/api_server/api_types.py b/src/xorl/server/api_server/api_types.py index d1c62786..6190dfb7 100644 --- a/src/xorl/server/api_server/api_types.py +++ b/src/xorl/server/api_server/api_types.py @@ -6,7 +6,7 @@ from typing import Any, Dict, List, Literal, Optional, Union -from pydantic import BaseModel, ConfigDict, Field, field_serializer, model_validator +from pydantic import AliasChoices, BaseModel, ConfigDict, Field, field_serializer, model_validator def _map_session_id_to_model_id(data: Any) -> Any: @@ -58,7 +58,12 @@ def tolist(self) -> List[Union[int, float, str]]: class Datum(BaseModel): - """Single training example with model inputs and loss function inputs.""" + """Single training example with model inputs and loss function inputs. + + `labels` with the same length as `input_ids` are treated as HF-format labels + and shifted for next-token prediction. Use `target_tokens` for already shifted + RL/xorl-client style targets. + """ model_input: Dict[str, InputType] = Field(..., description="Model input tensors (input_ids, position_ids, etc.)") loss_fn_inputs: Dict[str, InputType] = Field(..., description="Loss function input tensors (e.g., labels)") @@ -172,8 +177,25 @@ def model_dump_json(self, **kwargs): return super().model_dump_json(**kwargs) +class LoRAConfigRequest(BaseModel): + """Per-session LoRA overrides accepted by create_model.""" + + model_config = ConfigDict(extra="allow", populate_by_name=True) + + lora_rank: Optional[int] = Field( + default=None, + validation_alias=AliasChoices("lora_rank", "rank"), + description="LoRA rank override. Accepts Tinker's rank alias.", + ) + lora_alpha: Optional[int] = Field( + default=None, + validation_alias=AliasChoices("lora_alpha", "alpha"), + description="LoRA alpha override. Accepts Tinker's alpha alias.", + ) + + class AdamParams(BaseModel): - """AdamW optimizer parameters.""" + """Tinker-compatible AdamW optimizer parameters.""" learning_rate: float = Field(default=0.0001, description="Learning rate") beta1: float = Field(default=0.9, description="First moment coefficient") @@ -183,6 +205,41 @@ class AdamParams(BaseModel): grad_clip_norm: float = Field(default=0.0, description="Gradient clipping norm (0.0 = no clipping)") +class OptimizerConfigRequest(BaseModel): + """Per-session optimizer overrides accepted by create_model.""" + + model_config = ConfigDict(extra="allow") + + type: Optional[Literal["adamw", "anyprecision_adamw", "sgd", "signsgd", "muon"]] = Field( + default=None, description="Optimizer type" + ) + learning_rate: Optional[float] = Field(default=None, description="Default learning rate for the session") + weight_decay: Optional[float] = Field(default=None, description="Weight decay") + optimizer_dtype: Optional[Literal["fp32", "bf16"]] = Field(default=None, description="Optimizer state dtype") + betas: Optional[List[float]] = Field(default=None, description="Adam-family beta coefficients") + eps: Optional[float] = Field(default=None, description="Adam-family epsilon") + optimizer_kwargs: Optional[Dict[str, Any]] = Field(default=None, description="Optimizer-specific kwargs") + + +class LoRARuntimeConfig(BaseModel): + """Normalized LoRA runtime config returned by the API.""" + + lora_rank: int = Field(..., description="LoRA rank") + lora_alpha: int = Field(..., description="LoRA alpha") + + +class OptimizerRuntimeConfig(BaseModel): + """Normalized optimizer runtime config returned by the API.""" + + type: Literal["adamw", "anyprecision_adamw", "sgd", "signsgd", "muon"] = Field(..., description="Optimizer type") + learning_rate: float = Field(..., description="Default learning rate for the session") + weight_decay: float = Field(..., description="Weight decay") + optimizer_dtype: Literal["fp32", "bf16"] = Field(..., description="Optimizer state dtype") + betas: Optional[List[float]] = Field(default=None, description="Adam-family beta coefficients") + eps: Optional[float] = Field(default=None, description="Adam-family epsilon") + optimizer_kwargs: Dict[str, Any] = Field(default_factory=dict, description="Optimizer-specific kwargs") + + # ============================================================================ # Inference Operations # ============================================================================ @@ -265,14 +322,37 @@ class OptimStepRequest(BaseModel): default=None, description="Sequence ID for request ordering (ensures forward_backward executes before optim_step)", ) - adam_params: AdamParams = Field(default_factory=AdamParams, description="AdamW optimizer parameters") + learning_rate: Optional[float] = Field(default=None, description="Optional per-step learning rate override") gradient_clip: Optional[float] = Field(default=None, description="Gradient clipping value") + adam_params: Optional[AdamParams] = Field( + default=None, + description="Legacy Tinker AdamW optimizer parameters. learning_rate/gradient_clip take precedence.", + ) @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) + def _map_legacy_optimizer_fields(cls, data): + """Map legacy optimizer payloads onto the generic per-step LR field.""" + if isinstance(data, dict): + if "session_id" in data and "model_id" not in data: + data["model_id"] = data["session_id"] + if "lr" in data and "learning_rate" not in data: + data["learning_rate"] = data["lr"] + if "adam_params" in data: + adam_params = data.get("adam_params") or {} + if isinstance(adam_params, dict): + if "learning_rate" in adam_params and "learning_rate" not in data: + data["learning_rate"] = adam_params["learning_rate"] + if "grad_clip_norm" in adam_params and "gradient_clip" not in data: + data["gradient_clip"] = adam_params["grad_clip_norm"] + else: + learning_rate = getattr(adam_params, "learning_rate", None) + if learning_rate is not None and "learning_rate" not in data: + data["learning_rate"] = learning_rate + gradient_clip = getattr(adam_params, "grad_clip_norm", None) + if gradient_clip is not None and "gradient_clip" not in data: + data["gradient_clip"] = gradient_clip + return data class OptimStepResponse(BaseModel): @@ -371,6 +451,19 @@ class WeightsInfoResponse(BaseModel): base_model: str = Field(..., description="Base model name (e.g., 'Qwen/Qwen2.5-3B-Instruct')") is_lora: bool = Field(default=True, description="Whether this is a LoRA checkpoint") lora_rank: Optional[int] = Field(default=None, description="LoRA rank (if is_lora=True)") + lora_config: Optional[LoRARuntimeConfig] = Field( + default=None, description="Normalized LoRA runtime config for LoRA checkpoints" + ) + optimizer_config: Optional[OptimizerRuntimeConfig] = Field( + default=None, description="Normalized optimizer runtime config for LoRA checkpoints" + ) + + @model_validator(mode="after") + def _mirror_lora_rank(self): + """Keep Tinker's flat lora_rank field in sync with the richer metadata.""" + if self.lora_rank is None and self.lora_config is not None: + self.lora_rank = self.lora_config.lora_rank + return self class CreateModelRequest(BaseModel): @@ -380,19 +473,16 @@ class CreateModelRequest(BaseModel): model_id: str = Field(default="default", description="Model identifier") base_model: str = Field(..., description="Base model name (e.g., 'Qwen/Qwen2.5-3B-Instruct')") - lora_config: Dict[str, Any] = Field(default_factory=dict, description="LoRA configuration (rank, alpha, etc.)") + lora_config: Optional[LoRAConfigRequest] = Field(default=None, description="Per-session LoRA overrides") + optimizer_config: Optional[OptimizerConfigRequest] = Field( + default=None, description="Per-session optimizer configuration" + ) @model_validator(mode="before") @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): - # 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: - lora_cfg["lora_rank"] = lora_cfg["rank"] - return data + return _map_session_id_to_model_id(data) class CreateModelResponse(BaseModel): @@ -412,18 +502,15 @@ class CreateSessionRequest(BaseModel): 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") + lora_config: Optional[LoRAConfigRequest] = Field( + default=None, 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 + return _map_model_id_to_session_id(data) class CreateSessionResponse(BaseModel): @@ -652,11 +739,11 @@ class DeleteCheckpointResponse(BaseModel): class KillSessionRequest(BaseModel): - """API request for killing a full-weights training session. + """API request for killing a training session. - In full-weights training mode (enable_lora=False), the server operates in - single-tenant mode. This endpoint allows killing the active session to - start a new one. + In full-weight training mode, the server operates in single-tenant mode and + this endpoint clears the active session. In LoRA mode, it removes a + non-default tenant session from the worker and API registries. """ model_config = ConfigDict(extra="ignore") @@ -673,7 +760,7 @@ def _map_tinker_fields(cls, data): class KillSessionResponse(BaseModel): - """API response for killing a full-weights training session.""" + """API response for killing a training session.""" success: bool = Field(..., description="Whether the session was killed successfully") message: str = Field(..., description="Status message") @@ -732,6 +819,9 @@ class InferenceEndpoint(BaseModel): host: str = Field(..., description="Hostname or IP address of the inference endpoint") port: int = Field(..., description="Port number of the inference endpoint") + worker_port: Optional[int] = Field( + default=None, description="Port number of the inference worker endpoint, if different from port" + ) 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( @@ -744,6 +834,10 @@ class AddInferenceEndpointRequest(BaseModel): host: str = Field(..., description="Hostname or IP address of the inference endpoint") port: int = Field(..., description="Port number of the inference endpoint") + worker_port: Optional[int] = Field( + default=None, + description="Port number of the inference worker endpoint. Defaults to port when omitted.", + ) world_size: int = Field(default=1, description="Number of workers at this endpoint") # Auto-sync configuration sync_weights: bool = Field( @@ -809,7 +903,7 @@ class SyncInferenceWeightsRequest(BaseModel): "so stale KV entries from previous weights are evicted naturally.", ) pause_mode: Literal["retract", "abort", "in_place"] = Field( - default="retract", + default="in_place", description="How to pause inference during weight sync. " "'retract' (default): retract running requests to waiting queue, re-execute after resume. " "'abort': abort and return all in-flight requests. " @@ -844,6 +938,14 @@ class SyncInferenceWeightsResponse(BaseModel): total_bytes: int = Field(default=0, description="Total bytes transferred") num_parameters: int = Field(default=0, description="Number of parameters transferred") num_buckets: int = Field(default=0, description="Number of transfer buckets used") + timing_breakdown: Dict[str, float] = Field( + default_factory=dict, + description="Optional phase timing breakdown from the trainer handler, in seconds", + ) + p2p_rank_summaries: List[Dict[str, Any]] = Field( + default_factory=list, + description="Optional per-rank P2P transport timing summaries for tail-latency diagnosis", + ) endpoints_synced: List[EndpointSyncResult] = Field( default_factory=list, description="Sync results for each endpoint" ) @@ -888,10 +990,18 @@ class CreateSamplingSessionRequest(BaseModel): This loads the specified LoRA adapter on all inference workers. The model_path can be: + - 'xorl://model_id/sampler_weights/adapter_name' - 'sampler_weights/adapter_name' - 'adapter_name' (just the name) """ + model_id: Optional[str] = Field( + default=None, + description=( + "Training session to attribute the sampling adapter to for cleanup and LRU tracking. " + "When omitted, xorl:// model_id is used if present, otherwise 'default'." + ), + ) model_path: str = Field( ..., description="Path to the LoRA adapter (e.g., 'sampler_weights/adapter-001' or just 'adapter-001')", diff --git a/src/xorl/server/api_server/endpoints.py b/src/xorl/server/api_server/endpoints.py index f300bc7b..009c9edf 100644 --- a/src/xorl/server/api_server/endpoints.py +++ b/src/xorl/server/api_server/endpoints.py @@ -3,6 +3,7 @@ from __future__ import annotations import logging +import os import uuid from typing import Any, Dict, List, Optional @@ -53,7 +54,11 @@ ) 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 +from xorl.server.protocol.operations import KillSessionData, RegisterSessionData +from xorl.server.session_spec import ( + load_session_spec_from_checkpoint, + normalize_session_spec, +) logger = logging.getLogger(__name__) @@ -61,6 +66,104 @@ router = APIRouter() +def _dump_optional_config(config: Any) -> Optional[Dict[str, Any]]: + """Return a plain dict for optional Pydantic/dict config values.""" + if config is None: + return None + if hasattr(config, "model_dump"): + return config.model_dump(exclude_none=True) + return dict(config) + + +def _first_output_result(output: Any) -> Dict[str, Any]: + """Extract the first orchestrator output while tolerating dict/list shapes.""" + if isinstance(output.outputs, list): + return output.outputs[0] if output.outputs else {} + return output.outputs or {} + + +async def _register_runtime_session( + server, + *, + model_id: str, + base_model: str, + raw_lora_config: Optional[Dict[str, Any]], + raw_optimizer_config: Optional[Dict[str, Any]], +) -> Dict[str, Any]: + """Normalize and register a session runtime spec with workers and API state.""" + if server.base_model is not None and base_model != server.base_model: + raise ValueError( + f"create_model base_model must match the server base model. " + f"requested={base_model!r}, server={server.base_model!r}" + ) + + if server.default_session_spec is None: + if model_id != "default": + raise ValueError( + "Full-weight server multi-tenancy is not supported yet. Use the reserved model_id='default' session." + ) + if raw_lora_config or raw_optimizer_config: + raise ValueError("Per-session LoRA or optimizer overrides are not supported in full-weight server mode.") + normalized_spec = { + "base_model": base_model, + "is_lora": False, + } + materialize = False + else: + normalized_spec = normalize_session_spec( + base_model=base_model, + raw_lora_config=raw_lora_config, + raw_optimizer_config=raw_optimizer_config, + default_rank=server.default_session_spec["lora_config"]["lora_rank"], + default_alpha=server.default_session_spec["lora_config"]["lora_alpha"], + max_lora_rank=server.max_lora_rank or server.default_session_spec["lora_config"]["lora_rank"], + default_optimizer_type=server.default_session_spec["optimizer_config"]["type"], + default_learning_rate=server.default_session_spec["optimizer_config"]["learning_rate"], + default_weight_decay=server.default_session_spec["optimizer_config"]["weight_decay"], + default_optimizer_dtype=server.default_session_spec["optimizer_config"]["optimizer_dtype"], + default_optimizer_kwargs=server.default_session_spec["optimizer_config"].get("optimizer_kwargs", {}), + server_lora_config=server.server_lora_config, + default_betas=tuple(server.default_session_spec["optimizer_config"].get("betas") or (0.9, 0.95)), + default_eps=float(server.default_session_spec["optimizer_config"].get("eps") or 1e-8), + ) + materialize = True + + existing_spec = server.model_configs.get(model_id) + if existing_spec is not None: + if existing_spec != normalized_spec: + raise ValueError( + f"model_id={model_id!r} already exists with a different session spec. " + "Call /api/v1/unload_model first before recreating it." + ) + server._update_session_activity(model_id) + return normalized_spec + + engine_request = OrchestratorRequest( + operation="register_session", + payload=RegisterSessionData( + model_id=model_id, + session_spec=normalized_spec, + materialize=materialize, + ), + ) + response_future = await server.orchestrator_client.send_request(engine_request) + output = await server._wait_for_response( + response_future, + engine_request.request_id, + server.default_timeout, + "Register session timeout", + ) + result = _first_output_result(output) + register_result = result.get("result", result) + if not register_result.get("registered", False): + raise RuntimeError(register_result.get("error", f"Worker register_session failed for model_id={model_id}")) + + server.registered_model_ids.add(model_id) + server.model_configs[model_id] = normalized_spec + server._update_session_activity(model_id) + return normalized_spec + + # ============================================================================ # Training Endpoints (Two-Phase Pattern) # ============================================================================ @@ -139,7 +242,7 @@ async def forward_endpoint(request: ForwardRequest, server=Depends(require_api_s ) async def optim_step_endpoint(request: OptimStepRequest, server=Depends(require_api_server)): """ - Perform optimization step using AdamW optimizer (two-phase pattern). + Perform an optimizer step (two-phase pattern). Returns UntypedAPIFuture immediately. Poll /api/v1/retrieve_future to get the OptimStepResponse result. @@ -278,14 +381,12 @@ async def weights_info_endpoint(request: WeightsInfoRequest, server=Depends(requ """ Get checkpoint metadata for resuming training. - This endpoint returns the base_model and lora_rank needed to create - a TrainingClient that can load the checkpoint. It mirrors tinker's + This endpoint returns the full session runtime spec needed to recreate + a training session for a checkpoint. It mirrors tinker's /api/v1/weights_info endpoint for API compatibility. The xorl_path should be a xorl:// URI (e.g., "xorl://default/weights/checkpoint-001"). """ - # Parse the xorl:// URI to get model_id - # Format: xorl://model_id/weights/checkpoint_name xorl_path = request.xorl_path if not xorl_path.startswith("xorl://"): raise HTTPException( @@ -293,44 +394,46 @@ async def weights_info_endpoint(request: WeightsInfoRequest, server=Depends(requ detail=f"Invalid xorl_path format: {xorl_path}. Expected xorl://model_id/weights/checkpoint_name", ) - parts = xorl_path[7:].split("/") # Remove "xorl://" - if len(parts) < 1: + checkpoint_model_id, checkpoint_name, has_explicit_model_id = server._from_xorl_uri(xorl_path) + if not has_explicit_model_id or checkpoint_model_id is None: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, - detail=f"Invalid xorl_path format: {xorl_path}. Could not extract model_id", + detail=f"Invalid xorl_path format: {xorl_path}. Expected xorl://model_id/weights/checkpoint_name", ) + checkpoint_model_id = validate_model_id(checkpoint_model_id) - model_id = parts[0] - - # Look up model config - model_config = server.model_configs.get(model_id) - - if model_config is None: - # Fall back to server's base_model if no specific config is stored - if server.base_model is not None: - logger.warning( - f"No model config found for model_id '{model_id}', using server's base_model: {server.base_model}" - ) - return WeightsInfoResponse( - base_model=server.base_model, - is_lora=True, - lora_rank=None, # Unknown rank - ) - else: + weights_dir = os.path.abspath(os.path.join(server.output_dir, "weights", checkpoint_model_id)) + checkpoint_path = os.path.abspath(os.path.join(weights_dir, checkpoint_name)) + try: + if os.path.commonpath([checkpoint_path, weights_dir]) != weights_dir: raise HTTPException( - status_code=status.HTTP_404_NOT_FOUND, - detail=f"No model config found for model_id '{model_id}' and server has no default base_model configured", + status_code=status.HTTP_400_BAD_REQUEST, + detail=f"Invalid checkpoint path in xorl_path: {xorl_path}", ) + except ValueError: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail=f"Invalid checkpoint path in xorl_path: {xorl_path}", + ) + if not os.path.exists(checkpoint_path): + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=f"Checkpoint not found: {xorl_path}", + ) - # Extract lora_rank from lora_config - lora_config = model_config.get("lora_config", {}) - lora_rank = lora_config.get("lora_rank") + try: + session_spec = load_session_spec_from_checkpoint( + checkpoint_path, + fallback_base_model=server.base_model, + fallback_session_spec=server.model_configs.get(checkpoint_model_id) or server.default_session_spec, + ) + except Exception as e: + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail=f"Failed to read checkpoint metadata from {xorl_path}: {e}", + ) from e - return WeightsInfoResponse( - base_model=model_config["base_model"], - is_lora=True, - lora_rank=lora_rank, - ) + return WeightsInfoResponse(**session_spec) @router.post( @@ -350,54 +453,54 @@ async def create_model_endpoint(request: CreateModelRequest, server=Depends(requ Returns UntypedAPIFuture immediately. Poll /api/v1/retrieve_future to get the CreateModelResponse result. - This initializes the model on the training server but doesn't actually - do anything yet in the current implementation. It's a placeholder for - future multi-model support. - The base_model in the request must match the server's configured base model. """ if not server.future_store: raise HTTPException(status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail="Future store not initialized") - model_id = request.model_id + model_id = validate_model_id(request.model_id) + request_payload = request.model_dump() + request_payload["model_id"] = model_id + + async def ensure_reserved_checkpoint(model_id_to_save: str, *, overwrite_existing: bool = False) -> None: + """Ensure the reserved initial checkpoint exists for this session.""" + if getattr(server, "_skip_initial_checkpoint", False): + return + + checkpoint_path = os.path.join(server.output_dir, "weights", model_id_to_save, server.RESERVED_CHECKPOINT_NAME) + if os.path.exists(checkpoint_path) and not overwrite_existing: + return + + try: + save_request = SaveWeightsRequest( + model_id=model_id_to_save, + path=server.RESERVED_CHECKPOINT_NAME, + ) + save_response = await server.save_weights(save_request) + logger.info(f"Auto-saved initial checkpoint for model_id={model_id_to_save}: {save_response.path}") + except Exception as e: + logger.warning(f"Failed to auto-save initial checkpoint for model_id={model_id_to_save}: {e}") + # Don't fail create_model if checkpoint save fails - it's not critical async def process_create_model(request_data: Dict[str, Any]) -> Dict[str, Any]: """Process create_model request and return result dict.""" req = CreateModelRequest(**request_data) + req = req.model_copy(update={"model_id": validate_model_id(req.model_id)}) logger.info(f"Creating model: {req.model_id}, base_model={req.base_model}") - # Register the model_id so subsequent /api/v1/* calls can use it - server.registered_model_ids.add(req.model_id) - - # Store the model config for /api/v1/weights_info - server.model_configs[req.model_id] = { - "base_model": req.base_model, - "lora_config": req.lora_config, - } - - # Initialize session activity tracking - server._update_session_activity(req.model_id) + had_model_config = req.model_id in server.model_configs + normalized_spec = await _register_runtime_session( + server, + model_id=req.model_id, + base_model=req.base_model, + raw_lora_config=_dump_optional_config(req.lora_config), + raw_optimizer_config=_dump_optional_config(req.optimizer_config), + ) - logger.info(f"Registered model_id: {req.model_id} with lora_config: {req.lora_config}") + logger.info(f"Registered model_id: {req.model_id} with session_spec: {normalized_spec}") - # Auto-save initial checkpoint "000000" only once per server lifetime - # This preserves the initial model state (base LoRA weights) before any training - # Subsequent create_model calls skip this since the base model hasn't changed - if not server._initial_checkpoint_saved: - try: - save_request = SaveWeightsRequest( - model_id=req.model_id, - path=server.RESERVED_CHECKPOINT_NAME, # "000000" - ) - save_response = await server.save_weights(save_request) - server._initial_checkpoint_saved = True - logger.info(f"Auto-saved initial checkpoint for model_id={req.model_id}: {save_response.path}") - except Exception as e: - logger.warning(f"Failed to auto-save initial checkpoint for model_id={req.model_id}: {e}") - # Don't fail create_model if checkpoint save fails - it's not critical - else: - logger.debug(f"Skipping initial checkpoint save for model_id={req.model_id} (already saved)") + await ensure_reserved_checkpoint(req.model_id, overwrite_existing=not had_model_config) return CreateModelResponse(model_id=req.model_id).model_dump() @@ -406,7 +509,7 @@ async def process_create_model(request_data: Dict[str, Any]) -> Dict[str, Any]: model_id=model_id, request_type="create_model", process_fn=process_create_model, - request_data=request.model_dump(), + request_data=request_payload, ) return UntypedAPIFuture( @@ -449,6 +552,11 @@ async def unload_model_endpoint(request: UnloadModelRequest, server=Depends(requ raise HTTPException(status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail="Future store not initialized") model_id = request.model_id + if server.default_session_spec is not None and model_id == "default": + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail="The default LoRA session is reserved and cannot be unloaded.", + ) # Check if the session exists if model_id not in server.registered_model_ids: @@ -495,13 +603,14 @@ async def process_unload_model(request_data: Dict[str, Any]) -> Dict[str, Any]: ) async def kill_session_endpoint(request: KillSessionRequest, server=Depends(require_api_server)): """ - Kill the active full-weights training session. + Kill an active training session. In full-weights training mode (enable_lora=False), the server operates in single-tenant mode where only one training session is allowed at a time. This endpoint kills the active session to allow starting a new one. - For LoRA mode, this is a no-op since multi-tenancy is supported. + In LoRA mode, non-default sessions are removed from both worker state and + the API registry. The reserved "default" LoRA session is not removed. Args: request: KillSessionRequest with model_id to kill and save_checkpoint flag @@ -514,6 +623,13 @@ async def kill_session_endpoint(request: KillSessionRequest, server=Depends(requ logger.info(f"Killing session: {model_id}, save_checkpoint={save_checkpoint}") + if server.default_session_spec is not None and model_id == "default": + return KillSessionResponse( + success=True, + message="Default LoRA session is reserved and was not removed.", + checkpoint_path=None, + ) + try: engine_request = OrchestratorRequest( operation="kill_session", @@ -533,10 +649,24 @@ async def kill_session_endpoint(request: KillSessionRequest, server=Depends(requ result = output.outputs[0] if output.outputs else {} + if result.get("success", False) and server.default_session_spec is not None and model_id != "default": + await server._cleanup_session(model_id, notify_workers=False) + + checkpoint_path = result.get("checkpoint_path") + if checkpoint_path and server.default_session_spec is not None: + weights_dir = os.path.abspath(os.path.join(server.output_dir, "weights", model_id)) + checkpoint_abs = os.path.abspath(checkpoint_path) + try: + if os.path.commonpath([checkpoint_abs, weights_dir]) == weights_dir: + checkpoint_name = os.path.basename(os.path.normpath(checkpoint_abs)) + checkpoint_path = server._to_xorl_uri(model_id, checkpoint_name) + except ValueError: + pass + return KillSessionResponse( success=result.get("success", False), message=result.get("message", ""), - checkpoint_path=result.get("checkpoint_path"), + checkpoint_path=checkpoint_path, ) except HTTPException: @@ -808,23 +938,50 @@ async def create_session_endpoint( 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 + lora_config = request.lora_config.model_dump(exclude_none=True) if request.lora_config is not None else {} + + if server.default_session_spec is not None: + server._require_engine() + existing_spec = server.model_configs.get(session_id) + if existing_spec is not None and not request.lora_config: + if request.base_model is not None and request.base_model != existing_spec.get("base_model"): + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail=( + f"session_id={session_id!r} already exists with base_model={existing_spec.get('base_model')!r}; " + f"requested base_model={request.base_model!r}." + ), + ) + server._update_session_activity(session_id) + else: + base_model = request.base_model or server.base_model or server.default_session_spec["base_model"] + raw_lora_config = lora_config or None + try: + await _register_runtime_session( + server, + model_id=session_id, + base_model=base_model, + raw_lora_config=raw_lora_config, + raw_optimizer_config=None, + ) + except ValueError as e: + raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(e)) from e + else: + 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": lora_config, + }, + ) - 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 + if request.base_model: + model_config["base_model"] = request.base_model + if lora_config: + model_config["lora_config"] = lora_config - server._update_session_activity(session_id) + server._update_session_activity(session_id) warning_message = None info_message = f"Session '{session_id}' registered successfully." diff --git a/src/xorl/server/api_server/inference_endpoints.py b/src/xorl/server/api_server/inference_endpoints.py index a3156695..6283240b 100644 --- a/src/xorl/server/api_server/inference_endpoints.py +++ b/src/xorl/server/api_server/inference_endpoints.py @@ -11,6 +11,7 @@ from typing import Any, Dict, List import httpx +import requests from fastapi import HTTPException, status from huggingface_hub import hf_hub_download @@ -30,6 +31,7 @@ SyncInferenceWeightsRequest, SyncInferenceWeightsResponse, ) +from xorl.server.api_server.utils import validate_model_id from xorl.server.protocol.api_orchestrator import OrchestratorRequest from xorl.server.protocol.operations import SyncWeightsData @@ -40,6 +42,12 @@ class InferenceEndpointsMixin: """Mixin for inference endpoints, LoRA adapter management, and sampling sessions.""" + @staticmethod + def _endpoint_worker_url(endpoint: InferenceEndpoint) -> str: + """Return the inference worker URL used for LoRA adapter management.""" + worker_port = endpoint.worker_port if endpoint.worker_port is not None else endpoint.port + return f"http://{endpoint.host}:{worker_port}" + @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.""" @@ -149,6 +157,7 @@ async def _sync_weights_to_endpoints( master_port=master_port, group_name=group_name, buffer_size_mb=buffer_size_mb, + sync_method=self.sync_inference_method, quantization=quantization, ), ) @@ -189,6 +198,8 @@ 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 + worker_url = f"http://{request.host}:{worker_port}" # Check if endpoint already exists for existing in self.inference_endpoints: @@ -199,21 +210,27 @@ 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: - if not await self._check_endpoint_health(client, endpoint_url, "Inference endpoint"): + if not await self._check_endpoint_health(client, endpoint_url, "SGLang server"): raise Exception(f"All health endpoints failed for {endpoint_url}") + if worker_url != endpoint_url and not await self._check_endpoint_health( + client, worker_url, "Inference worker" + ): + 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}: {e}") + logger.warning(f"Health check failed for {endpoint_url} or {worker_url}: {e}") is_healthy = False if not is_healthy: return AddInferenceEndpointResponse( success=False, - message=f"Health check failed for inference endpoint {endpoint_url}", + message=f"Health check failed for SGLang server {endpoint_url} or inference worker {worker_url}", endpoint=None, ) @@ -294,6 +311,7 @@ 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, @@ -324,7 +342,7 @@ async def add_inference_endpoint(self, request: AddInferenceEndpointRequest) -> { "host": request.host, "port": request.port, - "world_size": request.world_size, + "world_size": world_size, } ], master_address=master_address, @@ -434,6 +452,9 @@ def remove_inference_endpoint(self, request: RemoveInferenceEndpointRequest) -> for i, endpoint in enumerate(self.inference_endpoints): if endpoint.host == request.host and endpoint.port == request.port: self.inference_endpoints.pop(i) + if not self.inference_endpoints and self.loaded_sampling_loras: + self.loaded_sampling_loras.clear() + logger.info("Cleared tracked sampling adapters after removing the last inference endpoint") logger.info(f"Removed inference endpoint: {endpoint_url}") return RemoveInferenceEndpointResponse( success=True, @@ -482,13 +503,7 @@ 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 @@ -551,6 +566,8 @@ async def sync_inference_weights(self, request: SyncInferenceWeightsRequest) -> total_bytes=result.get("total_bytes", 0), num_parameters=result.get("num_parameters", 0), num_buckets=result.get("num_buckets", 0), + timing_breakdown=result.get("timing_breakdown", {}), + p2p_rank_summaries=result.get("p2p_rank_summaries", []), endpoints_synced=endpoint_results, ) @@ -566,9 +583,9 @@ async def sync_inference_weights(self, request: SyncInferenceWeightsRequest) -> # Sampling Session Management (LoRA Adapter Loading) # ========================================================================= - def _resolve_model_path(self, model_path: str) -> tuple[str, str]: + def _resolve_model_path(self, model_path: str) -> tuple[str | None, str, str]: """ - Resolve model_path to (lora_name, absolute_path). + Resolve model_path to (source_model_id, lora_name, absolute_path). Sampler weights are stored flat under output_dir/sampler_weights/{name} without model_id subdirectories, because inference endpoints don't know about model_id. @@ -582,17 +599,19 @@ def _resolve_model_path(self, model_path: str) -> tuple[str, str]: model_path: Path to the LoRA adapter (can be xorl:// URI or relative path) Returns: - Tuple of (lora_name, absolute_path) + Tuple of (source_model_id, lora_name, absolute_path) Raises: HTTPException: If path format is invalid or path doesn't exist """ + source_model_id: str | None = None if model_path.startswith("xorl://"): # Format: xorl://model_id/sampler_weights/adapter_name # Parse: remove "xorl://", split by "/", extract adapter name parts = model_path[7:].split("/") # Remove "xorl://" if len(parts) >= 3 and parts[1] == "sampler_weights": # xorl://model_id/sampler_weights/adapter_name + source_model_id = validate_model_id(parts[0]) lora_name = "/".join(parts[2:]) # In case adapter name has / else: raise HTTPException( @@ -616,7 +635,7 @@ def _resolve_model_path(self, model_path: str) -> tuple[str, str]: status_code=status.HTTP_404_NOT_FOUND, detail=f"Model path does not exist: {absolute_path}" ) - return lora_name, absolute_path + return source_model_id, lora_name, absolute_path async def _load_lora_on_inference_endpoints(self, lora_name: str, lora_path: str) -> bool: """ @@ -639,37 +658,41 @@ async def _load_lora_on_inference_endpoints(self, lora_name: str, lora_path: str ) async def load_on_endpoint(endpoint: InferenceEndpoint) -> tuple[str, bool, str]: - endpoint_url = f"http://{endpoint.host}:{endpoint.port}" + endpoint_url = self._endpoint_worker_url(endpoint) try: - async with httpx.AsyncClient(timeout=300.0) as client: - response = await client.post( + + def do_post() -> requests.Response: + return requests.post( f"{endpoint_url}/load_lora_adapter", json={ "lora_name": lora_name, "lora_path": lora_path, "pinned": False, # Allow eviction for memory management }, + headers={"Connection": "close"}, + timeout=300.0, ) - result = response.json() + response = await asyncio.to_thread(do_post) + result = response.json() - if response.status_code == 200 and result.get("success", False): - logger.info(f"Loaded LoRA adapter '{lora_name}' on {endpoint_url}") - return endpoint_url, True, "" + if response.status_code == 200 and result.get("success", False): + logger.info(f"Loaded LoRA adapter '{lora_name}' on {endpoint_url}") + return endpoint_url, True, "" - # Check for errors - error_msg = result.get("error_message", "") + # Check for errors + error_msg = result.get("error_message", "") - # "already loaded" is not a fatal error - treat it as success - # This can happen when create_sampling_session is called after save_weights_for_sampler - if "already loaded" in error_msg.lower(): - logger.warning(f"LoRA adapter '{lora_name}' already loaded on {endpoint_url}, continuing") - return endpoint_url, True, "" + # "already loaded" is not a fatal error - treat it as success + # This can happen when create_sampling_session is called after save_weights_for_sampler + if "already loaded" in error_msg.lower(): + logger.warning(f"LoRA adapter '{lora_name}' already loaded on {endpoint_url}, continuing") + return endpoint_url, True, "" - if not error_msg: - error_msg = f"HTTP {response.status_code}" - logger.error(f"Failed to load LoRA adapter on {endpoint_url}: {error_msg}") - return endpoint_url, False, error_msg + if not error_msg: + error_msg = f"HTTP {response.status_code}" + logger.error(f"Failed to load LoRA adapter on {endpoint_url}: {error_msg}") + return endpoint_url, False, error_msg except Exception as e: logger.error(f"Failed to load LoRA adapter on {endpoint_url}: {e}") @@ -716,23 +739,28 @@ async def _unload_lora_on_inference_endpoints(self, lora_name: str) -> bool: return True async def unload_on_endpoint(endpoint: InferenceEndpoint) -> tuple[str, bool, str]: - endpoint_url = f"http://{endpoint.host}:{endpoint.port}" + endpoint_url = self._endpoint_worker_url(endpoint) try: - async with httpx.AsyncClient(timeout=30.0) as client: - response = await client.post( + + def do_post() -> requests.Response: + return requests.post( f"{endpoint_url}/unload_lora_adapter", json={"lora_name": lora_name}, + headers={"Connection": "close"}, + timeout=30.0, ) - response.raise_for_status() - result = response.json() - if result.get("success", False): - logger.info(f"Unloaded LoRA adapter '{lora_name}' from {endpoint_url}") - return endpoint_url, True, "" - else: - error_msg = result.get("error_message", "Unknown error") - logger.warning(f"Failed to unload LoRA adapter from {endpoint_url}: {error_msg}") - return endpoint_url, False, error_msg + response = await asyncio.to_thread(do_post) + response.raise_for_status() + result = response.json() + + if result.get("success", False): + logger.info(f"Unloaded LoRA adapter '{lora_name}' from {endpoint_url}") + return endpoint_url, True, "" + else: + error_msg = result.get("error_message", "Unknown error") + logger.warning(f"Failed to unload LoRA adapter from {endpoint_url}: {error_msg}") + return endpoint_url, False, error_msg except Exception as e: logger.warning(f"Failed to unload LoRA adapter from {endpoint_url}: {e}") @@ -750,7 +778,7 @@ async def unload_on_endpoint(endpoint: InferenceEndpoint) -> tuple[str, bool, st return True - async def _get_loaded_adapters_from_endpoint(self, endpoint: "InferenceEndpoint") -> list[str]: + async def _get_loaded_adapters_from_endpoint(self, endpoint: "InferenceEndpoint") -> list[str] | None: """ Get list of currently loaded LoRA adapters from an inference endpoint. @@ -758,9 +786,9 @@ async def _get_loaded_adapters_from_endpoint(self, endpoint: "InferenceEndpoint" endpoint: The inference endpoint to query Returns: - List of adapter names currently loaded + List of adapter names currently loaded, or None if the endpoint could not be queried """ - endpoint_url = f"http://{endpoint.host}:{endpoint.port}" + endpoint_url = self._endpoint_worker_url(endpoint) try: async with httpx.AsyncClient(timeout=10.0) as client: response = await client.get(f"{endpoint_url}/v1/models") @@ -775,8 +803,57 @@ async def _get_loaded_adapters_from_endpoint(self, endpoint: "InferenceEndpoint" return adapters except Exception as e: logger.warning(f"Error getting loaded adapters from {endpoint_url}: {e}") + return None + + async def _reconcile_tracked_adapters(self, model_id: str) -> list[set[str]]: + """ + Reconcile tracked adapter state with what inference endpoints actually have loaded. + + This prunes stale tracking entries left behind by endpoint restarts or failed + create_sampling_session attempts, while preserving adapters that are still loaded + on at least one endpoint and may just need reloading on the others. + + Returns: + Per-endpoint sets of currently loaded adapter names. + """ + if not self.inference_endpoints: return [] + results = await asyncio.gather( + *[self._get_loaded_adapters_from_endpoint(endpoint) for endpoint in self.inference_endpoints], + return_exceptions=True, + ) + + loaded_by_endpoint: list[set[str]] = [] + unknown_queries = False + for result in results: + if isinstance(result, Exception): + logger.warning(f"Exception while querying loaded adapters: {result}") + unknown_queries = True + loaded_by_endpoint.append(set()) + elif result is None: + unknown_queries = True + loaded_by_endpoint.append(set()) + else: + loaded_by_endpoint.append(set(result)) + + loaded_anywhere = set().union(*loaded_by_endpoint) if loaded_by_endpoint else set() + tracked = self.loaded_sampling_loras.get(model_id, []) + if tracked and not unknown_queries: + stale = [name for name, _ in tracked if name not in loaded_anywhere] + if stale: + self.loaded_sampling_loras[model_id] = [ + (name, path) for name, path in tracked if name in loaded_anywhere + ] + logger.info(f"Pruned {len(stale)} stale tracked adapter(s) for model_id={model_id}: {stale}") + elif tracked: + logger.info( + f"Skipped stale adapter pruning for model_id={model_id} because one or more endpoints " + "could not report loaded adapters" + ) + + return loaded_by_endpoint + async def _unload_all_adapters_from_endpoints(self) -> int: """ Unload all currently loaded LoRA adapters from all inference endpoints. @@ -832,20 +909,29 @@ async def _unload_adapters_for_model(self, model_id: str) -> int: logger.info(f"Unloaded {unloaded} adapter(s) for model_id={model_id}") return unloaded - def _track_adapter(self, lora_name: str, lora_path: str, model_id: str = "default") -> bool: + def _track_adapter( + self, + lora_name: str, + lora_path: str, + model_id: str = "default", + *, + add_if_missing: bool = True, + ) -> bool: """ Track a loaded adapter in the LRU list. If the adapter is already tracked, moves it to MRU position. - If not tracked, adds it to the tracking list. + If not tracked, optionally adds it to the tracking list. Args: lora_name: Name of the LoRA adapter lora_path: Path to the adapter files model_id: The model/session ID for per-session tracking (default: "default") + add_if_missing: When False, only touches existing entries and does not create + a new tracking entry. Returns: - True if adapter was already tracked, False if newly added + True if adapter was already tracked, False otherwise """ # Initialize tracking list if not exists @@ -859,10 +945,13 @@ def _track_adapter(self, lora_name: str, lora_path: str, model_id: str = "defaul if existing_name == lora_name: # Move to end (most recently used) adapters.remove((existing_name, existing_path)) - adapters.append((existing_name, existing_path)) + adapters.append((lora_name, lora_path)) logger.info(f"LoRA adapter '{lora_name}' already tracked, moved to MRU") return True + if not add_if_missing: + return False + # Add new adapter to tracking adapters.append((lora_name, lora_path)) logger.info(f"LoRA adapter '{lora_name}' added to tracking list (count={len(adapters)})") @@ -885,32 +974,45 @@ async def create_sampling_session(self, request: CreateSamplingSessionRequest) - Response with session info """ model_path = request.model_path - model_id = getattr(request, "model_id", "default") or "default" + requested_model_id = validate_model_id(request.model_id) logger.info(f"Creating sampling session for model_path: {model_path}") # Resolve path and validate - lora_name, absolute_path = self._resolve_model_path(model_path) + path_model_id, lora_name, absolute_path = self._resolve_model_path(model_path) + model_id = path_model_id or requested_model_id + logger.info(f"Sampling session will be tracked under model_id={model_id}") + + loaded_by_endpoint = await self._reconcile_tracked_adapters(model_id) + loaded_on_all_endpoints = bool(loaded_by_endpoint) and all( + lora_name in loaded_names for loaded_names in loaded_by_endpoint + ) - # Check if already tracked (returns True if already tracked and moves to MRU) - already_tracked = self._track_adapter(lora_name, absolute_path, model_id=model_id) + # Touch existing tracking entry if present, but do not create a new entry until we + # know the adapter is actually loaded or the load succeeds. + already_tracked = self._track_adapter( + lora_name, + absolute_path, + model_id=model_id, + add_if_missing=False, + ) - if already_tracked: - # Already loaded by save_weights_for_sampler, nothing to do + if loaded_on_all_endpoints: + if not already_tracked: + self._track_adapter(lora_name, absolute_path, model_id=model_id) logger.info(f"LoRA adapter '{lora_name}' already loaded, skipping duplicate load") else: - # Need to load - but first check if we need to evict oldest (LRU) adapters = self.loaded_sampling_loras.get(model_id, []) - 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] + if not already_tracked and len(adapters) >= self.max_adapters_per_model: + oldest_name, _oldest_path = adapters[0] 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) - # Load on inference endpoints await self._load_lora_on_inference_endpoints(lora_name, absolute_path) + if not already_tracked: + self._track_adapter(lora_name, absolute_path, model_id=model_id) total_adapters = len(self.loaded_sampling_loras.get(model_id, [])) logger.info(f"Sampling session created: lora_name={lora_name}, total_adapters={total_adapters}") diff --git a/src/xorl/server/api_server/orchestrator_client.py b/src/xorl/server/api_server/orchestrator_client.py index 993c7599..fe15e31e 100644 --- a/src/xorl/server/api_server/orchestrator_client.py +++ b/src/xorl/server/api_server/orchestrator_client.py @@ -5,7 +5,7 @@ Architecture: - INPUT SOCKET: ZMQ ROUTER (bind) - sends requests to engine -- OUTPUT SOCKET: ZMQ PULL (connect) - receives outputs from engine +- OUTPUT SOCKET: ZMQ PULL (bind) - receives outputs from engine - Background task: process_outputs_socket() - continuously reads from PULL socket - Asyncio queue: output_queue - buffers outputs for async retrieval @@ -60,7 +60,7 @@ def __init__( Args: input_addr: ZMQ address for input ROUTER socket (e.g., "tcp://127.0.0.1:5555") - output_addr: ZMQ address for output PULL socket (e.g., "tcp://127.0.0.1:5556") + output_addr: ZMQ address for output PULL socket to bind (e.g., "tcp://127.0.0.1:5556") output_queue_maxsize: Maximum size of output queue """ self.input_addr = input_addr @@ -105,9 +105,11 @@ async def start(self): self.input_channel = AsyncRouterChannel(self.input_addr, context=self.context) self.input_channel.bind() - # Create OUTPUT channel (PULL for receiving outputs) + # Create OUTPUT channel (PULL for receiving outputs). The API process + # starts before the orchestrator is ready to emit responses, so bind + # the consumer side here and let the orchestrator's PUSH connect later. self.output_channel = AsyncPullChannel(self.output_addr, context=self.context) - self.output_channel.connect() + self.output_channel.bind() # Create output queue self.output_queue = asyncio.Queue(maxsize=self.output_queue_maxsize) diff --git a/src/xorl/server/api_server/server.py b/src/xorl/server/api_server/server.py index 2ee23132..a973ea82 100644 --- a/src/xorl/server/api_server/server.py +++ b/src/xorl/server/api_server/server.py @@ -44,6 +44,7 @@ import uvicorn from fastapi import FastAPI, HTTPException, status +import xorl.server.api_server._state as _state from xorl.server.api_server.api_types import ( InferenceEndpoint, LossFnOutput, @@ -61,6 +62,7 @@ from xorl.server.api_server.weights import WeightsMixin from xorl.server.protocol.api_orchestrator import OrchestratorRequest from xorl.server.protocol.operations import KillSessionData +from xorl.server.session_spec import normalize_session_spec logger = logging.getLogger(__name__) @@ -86,18 +88,23 @@ def __init__( default_timeout: float = 120.0, output_dir: str = "outputs", base_model: Optional[str] = None, + default_session_spec: Optional[Dict[str, Any]] = None, + server_lora_config: Optional[Dict[str, Any]] = None, + max_lora_rank: Optional[int] = None, storage_limit: str = "10TB", max_sampling_loras: int = 3, idle_session_timeout: float = 7200.0, skip_initial_checkpoint: bool = False, sync_inference_method: str = "nccl_broadcast", + train_config: Optional[Dict[str, Any]] = None, + lora_config: Optional[Dict[str, Any]] = None, ): """ Initialize Unified API Server. Args: engine_input_addr: Engine input address (ROUTER binds here) - engine_output_addr: Engine output address (PULL connects here) + engine_output_addr: Engine output address (PULL binds here) default_timeout: Default timeout for engine operations output_dir: Output directory for checkpoints and sampler weights (must be on shared filesystem) base_model: Base model name that this server is configured for (e.g., 'Qwen/Qwen2.5-3B-Instruct'). @@ -112,15 +119,25 @@ def __init__( This is useful for full-weight mode to avoid memory issues during save. sync_inference_method: Method for syncing weights to inference endpoints. Currently only 'nccl_broadcast' is supported. + train_config: Server train config used as defaults for per-session optimizer specs. + lora_config: Server LoRA config used as defaults and limits for per-session LoRA specs. """ self.engine_input_addr = engine_input_addr self.engine_output_addr = engine_output_addr self.default_timeout = default_timeout self.output_dir = output_dir self.base_model = base_model + self.default_session_spec = default_session_spec + self.server_lora_config = server_lora_config or {} + self.max_lora_rank = max_lora_rank or self.server_lora_config.get( + "max_lora_rank", + self.default_session_spec["lora_config"]["lora_rank"] if self.default_session_spec is not None else None, + ) self.storage_limit = storage_limit self.max_sampling_loras = max_sampling_loras self.sync_inference_method = sync_inference_method + self.train_config = dict(train_config or {}) + self.lora_config = dict(lora_config or {}) # OrchestratorClient self.orchestrator_client: Optional[OrchestratorClient] = None @@ -143,18 +160,22 @@ def __init__( # "default" is pre-registered to allow direct API usage without create_model self.registered_model_ids: set = {"default"} - # Model config registry for storing LoRA configs - # Maps model_id -> {"base_model": str, "lora_config": dict} - # Used by /api/v1/weights_info to return checkpoint metadata + # Normalized session-spec registry + # Maps model_id -> {"base_model", "lora_config", "optimizer_config", ...} self.model_configs: Dict[str, Dict[str, Any]] = {} + if default_session_spec is not None: + self.model_configs["default"] = default_session_spec # Sampling session LoRA tracking (per-model_id, LRU order - oldest first) # Maps model_id -> List of (lora_name, model_path) tuples for loaded adapters # Each model_id has its own set of tracked adapters to support parallel training runs self.loaded_sampling_loras: Dict[str, List[tuple]] = {} - # Maximum number of adapters per model_id for sampling - self.max_adapters_per_model: int = 3 + # Maximum number of adapters per model_id for sampling. Default is intentionally + # generous (was 3) because SGLang's /unload_lora_adapter has been observed to + # hang on 30B-class hosts, so eviction during a multi-step OPD run wedges the + # session. Override via XORL_MAX_ADAPTERS_PER_MODEL when needed. + self.max_adapters_per_model: int = int(os.environ.get("XORL_MAX_ADAPTERS_PER_MODEL", "32")) # Session activity tracking for idle cleanup # Maps model_id -> last activity timestamp (time.time()) @@ -171,10 +192,9 @@ def __init__( # Organized by model_id for session-based cleanup self.future_store: Optional[FutureStore] = None - # Flag to track if initial checkpoint "000000" has been saved - # This should only happen once when the first model is created after server start - # Set to True if skip_initial_checkpoint is True to prevent auto-save in create_model - self._initial_checkpoint_saved: bool = skip_initial_checkpoint + # Whether create_model should skip auto-saving the reserved initial + # checkpoint "000000" for each training session. + self._skip_initial_checkpoint: bool = skip_initial_checkpoint logger.info( f"APIServer initialized: " @@ -183,6 +203,32 @@ def __init__( f"idle_session_timeout={self.idle_session_timeout}s" ) + def build_lora_session_spec( + self, + *, + base_model: str, + raw_lora_config: Optional[Dict[str, Any]] = None, + raw_optimizer_config: Optional[Dict[str, Any]] = None, + ) -> Dict[str, Any]: + """Build the normalized worker session spec for a create_model request.""" + server_lora_config = dict(self.lora_config or {}) + default_rank = int(server_lora_config.get("lora_rank", 32)) + max_lora_rank = int(server_lora_config.get("max_lora_rank", default_rank)) + return normalize_session_spec( + base_model=base_model, + raw_lora_config=raw_lora_config, + raw_optimizer_config=raw_optimizer_config, + default_rank=default_rank, + default_alpha=int(server_lora_config.get("lora_alpha", 16)), + max_lora_rank=max_lora_rank, + default_optimizer_type=self.train_config.get("optimizer", "adamw"), + default_learning_rate=float(self.train_config.get("lr", 1e-5)), + default_weight_decay=float(self.train_config.get("weight_decay", 0.01)), + default_optimizer_dtype=self.train_config.get("optimizer_dtype", "bf16"), + default_optimizer_kwargs=self.train_config.get("optimizer_kwargs", {}), + server_lora_config=server_lora_config, + ) + def validate_model_id(self, model_id: str) -> None: """ Validate that a model_id has been registered via create_model. @@ -265,11 +311,15 @@ def _build_loss_fn_outputs(result: Dict[str, Any]): @staticmethod def _build_info(result: Dict[str, Any]) -> Dict[str, Any]: """Extract auto-load info from engine result.""" + info: Dict[str, Any] = {} if result.get("auto_loaded"): - return {"auto_loaded": True, "auto_load_path": result.get("auto_load_path")} - return {} + info["auto_loaded"] = True + info["auto_load_path"] = result.get("auto_load_path") + if "teacher_hidden_cache" in result: + info["teacher_hidden_cache"] = result["teacher_hidden_cache"] + return info - async def _cleanup_session(self, model_id: str) -> None: + async def _cleanup_session(self, model_id: str, *, notify_workers: bool = True) -> None: """ Clean up all server-side state for a session. @@ -279,37 +329,41 @@ async def _cleanup_session(self, model_id: str) -> None: Args: model_id: The model identifier to clean up + notify_workers: Whether to send a kill_session command to workers first """ logger.info(f"Cleaning up session: {model_id}") + preserve_default_registration = model_id == "default" - # For full-weights training mode, send kill_session to workers to reset their state - # This ensures workers don't reject new sessions due to stale active session - try: - engine_request = OrchestratorRequest( - operation="kill_session", - payload=KillSessionData( - model_id=model_id, - save_checkpoint=False, # Don't save checkpoint on idle cleanup - ), - ) + if notify_workers: + # Send kill_session to workers before dropping local state. + try: + engine_request = OrchestratorRequest( + operation="kill_session", + payload=KillSessionData( + model_id=model_id, + save_checkpoint=False, # Don't save checkpoint on idle cleanup + ), + ) - response_future = await self.orchestrator_client.send_request(engine_request) - output = await self._wait_for_response( - response_future, - engine_request.request_id, - timeout=60.0, # Shorter timeout for cleanup - timeout_message=f"Kill session timeout during cleanup for {model_id}", - ) + response_future = await self.orchestrator_client.send_request(engine_request) + output = await self._wait_for_response( + response_future, + engine_request.request_id, + timeout=60.0, # Shorter timeout for cleanup + timeout_message=f"Kill session timeout during cleanup for {model_id}", + ) - result = output.outputs[0] if output.outputs else {} - if result.get("success"): - logger.info(f"Workers acknowledged session cleanup for {model_id}") - else: - logger.warning(f"Workers returned non-success for session cleanup: {result.get('message', 'unknown')}") + result = output.outputs[0] if output.outputs else {} + if result.get("success"): + logger.info(f"Workers acknowledged session cleanup for {model_id}") + else: + logger.warning( + f"Workers returned non-success for session cleanup: {result.get('message', 'unknown')}" + ) - except Exception as e: - # Log but don't fail - we still want to clean up local state - logger.warning(f"Failed to notify workers of session cleanup for {model_id}: {e}") + except Exception as e: + # Log but don't fail - we still want to clean up local state + logger.warning(f"Failed to notify workers of session cleanup for {model_id}: {e}") # Unload sampling adapters from SGL inference endpoints BEFORE removing tracking # This ensures we actually send unload requests to SGL @@ -323,8 +377,9 @@ async def _cleanup_session(self, model_id: str) -> None: # Remove from tracking structures # Note: _unload_adapters_for_model clears loaded_sampling_loras[model_id] to [], # so we pop to fully remove the entry - self.registered_model_ids.discard(model_id) - self.model_configs.pop(model_id, None) + if not preserve_default_registration: + self.registered_model_ids.discard(model_id) + self.model_configs.pop(model_id, None) self.session_last_activity.pop(model_id, None) self.loaded_sampling_loras.pop(model_id, None) @@ -334,9 +389,6 @@ async def _cleanup_session(self, model_id: str) -> None: if deleted_count > 0: logger.info(f"Cleaned up {deleted_count} futures for session {model_id}") - # Training adapters on workers are managed via LRU eviction - # No explicit unload needed - workers will evict when memory pressure occurs - logger.info(f"Session {model_id} cleaned up successfully") async def _cleanup_idle_sessions(self) -> None: @@ -357,6 +409,7 @@ async def _cleanup_idle_sessions(self) -> None: (model_id, last_activity) for model_id, last_activity in list(self.session_last_activity.items()) if current_time - last_activity > self.idle_session_timeout + and not (self.default_session_spec is not None and model_id == "default") ] for model_id, last_activity in idle_model_ids: @@ -462,13 +515,6 @@ async def stop(self): logger.info("APIServer stopped") -# ============================================================================ -# Global API Server Instance -# ============================================================================ - -import xorl.server.api_server._state as _state - - @asynccontextmanager async def lifespan(app: FastAPI): """Lifecycle manager for FastAPI app.""" diff --git a/src/xorl/server/api_server/training_ops.py b/src/xorl/server/api_server/training_ops.py index 76b1f1a0..b2048cff 100644 --- a/src/xorl/server/api_server/training_ops.py +++ b/src/xorl/server/api_server/training_ops.py @@ -5,11 +5,12 @@ import logging import math import time -from typing import Any, Dict +from typing import Any, Dict, Optional from fastapi import HTTPException, status from xorl.server.api_server.api_types import ( + AdamParams, ForwardBackwardRequest, ForwardBackwardResponse, ForwardRequest, @@ -51,6 +52,60 @@ def _sanitize_nan_to_zero(data): class TrainingOpsMixin: """Mixin for two-phase async pattern and core training operations.""" + def _session_default_learning_rate(self, model_id: str) -> Optional[float]: + """Return the registered session's default optimizer learning rate, if any.""" + model_configs = getattr(self, "model_configs", {}) + model_config = model_configs.get(model_id) or {} + optimizer_config = model_config.get("optimizer_config") or {} + if isinstance(optimizer_config, dict): + learning_rate = optimizer_config.get("learning_rate", optimizer_config.get("lr")) + if learning_rate is not None: + return float(learning_rate) + return None + + def _server_default_learning_rate(self) -> Optional[float]: + """Return the server train-config learning rate for full-weight sessions.""" + train_config = getattr(self, "train_config", {}) or {} + if not isinstance(train_config, dict): + return None + learning_rate = train_config.get("learning_rate", train_config.get("lr")) + return float(learning_rate) if learning_rate is not None else None + + def _optim_step_learning_rate(self, request: OptimStepRequest) -> float: + """Resolve the effective LR for an optim_step request. + + Priority: request.learning_rate, request.adam_params.learning_rate, + per-session optimizer_config, server train_config. If the session was + explicitly registered (even without an optimizer_config) fall back to + AdamParams().learning_rate; otherwise raise so a missing LR fails loud. + """ + if getattr(request, "learning_rate", None) is not None: + return float(request.learning_rate) + + fields_set = getattr(request, "model_fields_set", set()) + adam_params = getattr(request, "adam_params", None) + if "adam_params" in fields_set and adam_params is not None and adam_params.learning_rate is not None: + return float(adam_params.learning_rate) + + session_lr = self._session_default_learning_rate(request.model_id) + if session_lr is not None: + return session_lr + + server_lr = self._server_default_learning_rate() + if server_lr is not None: + return server_lr + + if request.model_id in getattr(self, "model_configs", {}): + return AdamParams().learning_rate + + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail=( + "optim_step: no learning_rate in request, no default optimizer_config " + f"registered for model_id={request.model_id!r}, and no server train_config lr" + ), + ) + # ========================================================================= # Two-Phase Request Pattern Methods # ========================================================================= @@ -245,11 +300,11 @@ async def forward_backward(self, request: ForwardBackwardRequest) -> ForwardBack # Debug: Log what we got from the engine logger.debug(f"API Server: Received result from engine, keys: {list(result.keys())}") - is_metrics = {k: v for k, v in result.items() if k.startswith("is_")} - if is_metrics: - logger.debug(f"API Server: IS metrics present in result: {list(is_metrics.keys())}") + loss_metrics = {k: v for k, v in result.items() if k.startswith(("is_", "opd_"))} + if loss_metrics: + logger.debug(f"API Server: loss metrics present in result: {list(loss_metrics.keys())}") else: - logger.debug("API Server: No IS metrics in result") + logger.debug("API Server: No loss metrics in result") # Sanitize NaN/Inf values for JSON serialization result = _sanitize_nan_to_zero(result) @@ -263,12 +318,12 @@ async def forward_backward(self, request: ForwardBackwardRequest) -> ForwardBack "loss:sum": total_loss * valid_tokens, "loss:mean": total_loss, "valid_tokens:sum": valid_tokens, - "execution_time:sum": result.get("execution_time", 0.0), + "execution_time:sum": result.get("execution_time", result.get("forward_backward_time", 0.0)), } - # Add IS metrics if present (already have name:reduction format) + # Add loss-specific metrics if present (already have name:reduction format) for key, value in result.items(): - if key.startswith("is_"): + if key.startswith(("is_", "opd_")): # Ensure colon format for tinker compatibility metrics[key if ":" in key else f"{key}:mean"] = value @@ -356,12 +411,21 @@ async def forward(self, request: ForwardRequest) -> ForwardResponse: "valid_tokens": valid_tokens, "execution_time": result.get("execution_time", 0.0), } + for key, value in result.items(): + if key.startswith(("is_", "opd_")): + metrics[key if ":" in key else f"{key}:mean"] = value + elif key in ( + "teacher_prefill_tokens", + "teacher_prefill_forward_compute_s", + "teacher_hidden_cache_write_s", + ): + metrics[key] = value return ForwardResponse( loss_fn_output_type=loss_fn_output_type, loss_fn_outputs=loss_fn_outputs, metrics=metrics, - info={}, + info=self._build_info(result), ) except HTTPException: @@ -387,19 +451,25 @@ async def optim_step(self, request: OptimStepRequest) -> OptimStepResponse: self._require_engine() try: + adam_params = request.adam_params + lr = self._optim_step_learning_rate(request) + # Determine gradient clipping value # Priority: explicit gradient_clip parameter, then adam_params.grad_clip_norm gradient_clip = request.gradient_clip - if gradient_clip is None and request.adam_params.grad_clip_norm > 0: - gradient_clip = request.adam_params.grad_clip_norm + if gradient_clip is None and adam_params is not None and adam_params.grad_clip_norm > 0: + gradient_clip = adam_params.grad_clip_norm # Create engine request # Pass seq_id and model_id for request ordering (SeqIdAwareFIFOPolicy) engine_request = OrchestratorRequest( operation="optim_step", payload=OptimStepData( - lr=request.adam_params.learning_rate, + lr=lr, gradient_clip=gradient_clip, + beta1=adam_params.beta1 if adam_params is not None else None, + beta2=adam_params.beta2 if adam_params is not None else None, + eps=adam_params.eps if adam_params is not None else None, model_id=request.model_id, ), seq_id=request.seq_id, @@ -435,11 +505,12 @@ async def optim_step(self, request: OptimStepRequest) -> OptimStepResponse: ) grad_norm = _sanitize_nan_to_zero(result.get("grad_norm", 0.0)) + response_learning_rate = result.get("learning_rate", result.get("lr", lr)) return OptimStepResponse( metrics={ "grad_norm": grad_norm, - "learning_rate": result.get("lr", request.adam_params.learning_rate), + "learning_rate": response_learning_rate, }, info=info, ) diff --git a/src/xorl/server/api_server/weights.py b/src/xorl/server/api_server/weights.py index cb69175f..979d0b06 100644 --- a/src/xorl/server/api_server/weights.py +++ b/src/xorl/server/api_server/weights.py @@ -39,6 +39,19 @@ logger = logging.getLogger(__name__) +def _model_config_is_lora(model_config: dict, *, default: bool = False) -> bool: + """Infer whether a stored model config represents a LoRA session.""" + if "is_lora" in model_config: + return bool(model_config["is_lora"]) + + lora_config = model_config.get("lora_config") or {} + if lora_config.get("enable_lora", False): + return True + if "lora_rank" in lora_config or "rank" in lora_config: + return True + return default + + class WeightsMixin: """Mixin for weight I/O and checkpoint management.""" @@ -517,6 +530,7 @@ def list_training_runs( base_model = model_config.get("base_model", self.base_model or "unknown") lora_config = model_config.get("lora_config", {}) lora_rank = lora_config.get("lora_rank") + is_lora = _model_config_is_lora(model_config, default=self.default_session_spec is not None) # Get last checkpoint info last_checkpoint = None @@ -537,7 +551,7 @@ def list_training_runs( training_run_id=model_id, base_model=base_model, model_owner="local", - is_lora=True, + is_lora=is_lora, corrupted=False, lora_rank=lora_rank, last_request_time=datetime.now().isoformat(), @@ -700,9 +714,18 @@ 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) or "rank" in lora_config - merge_lora_interval = lora_config.get("merge_lora_interval", 0) + lora_config = model_config.get("lora_config") or {} + is_lora = _model_config_is_lora(model_config, default=self.default_session_spec is not None) + if is_lora: + merge_lora_interval = int( + (getattr(self, "server_lora_config", {}) or {}).get( + "merge_lora_interval", + lora_config.get("merge_lora_interval", 0), + ) + or 0 + ) + else: + merge_lora_interval = 0 if is_lora and merge_lora_interval == 0: # LoRA with no merge: base weights unchanged, save adapter only diff --git a/src/xorl/server/backend/base.py b/src/xorl/server/backend/base.py index 7f195f44..ae8503d9 100644 --- a/src/xorl/server/backend/base.py +++ b/src/xorl/server/backend/base.py @@ -28,6 +28,7 @@ async def forward_backward( loss_fn_params: Optional[Dict[str, Any]] = None, model_id: Optional[str] = None, routed_experts: Optional[List[Any]] = None, + routed_expert_logits: Optional[List[Any]] = None, request_id: Optional[str] = None, ) -> Dict[str, Any]: """Run forward + backward pass. Returns {total_loss, global_valid_tokens, ...}.""" @@ -39,6 +40,8 @@ async def forward( loss_fn: str = "causallm_loss", loss_fn_params: Optional[Dict[str, Any]] = None, model_id: Optional[str] = None, + routed_experts: Optional[List[Any]] = None, + routed_expert_logits: Optional[List[Any]] = None, request_id: Optional[str] = None, ) -> Dict[str, Any]: """Forward-only pass (no gradients). Returns {total_loss, global_valid_tokens, ...}.""" @@ -136,6 +139,16 @@ async def sync_inference_weights( # Adapter Operations # ======================================================================== + @abstractmethod + async def register_session( + self, + model_id: str = "default", + session_spec: Optional[Dict[str, Any]] = None, + materialize: bool = False, + request_id: Optional[str] = None, + ) -> Dict[str, Any]: + """Register a session runtime spec on workers.""" + @abstractmethod async def register_adapter( self, diff --git a/src/xorl/server/backend/dummy.py b/src/xorl/server/backend/dummy.py index ecbba254..8e960f29 100644 --- a/src/xorl/server/backend/dummy.py +++ b/src/xorl/server/backend/dummy.py @@ -39,7 +39,14 @@ def _maybe_fail(self, operation: str): 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 + self, + batches, + loss_fn="causallm_loss", + loss_fn_params=None, + model_id=None, + routed_experts=None, + routed_expert_logits=None, + request_id=None, ): self._maybe_fail("forward_backward") valid_tokens = sum(len(batch.get("input_ids", [])) for batch in (batches or [])) @@ -49,7 +56,16 @@ async def forward_backward( "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, + routed_experts=None, + routed_expert_logits=None, + request_id=None, + ): self._maybe_fail("forward") valid_tokens = sum(len(batch.get("input_ids", [])) for batch in (batches or [])) return { @@ -108,6 +124,14 @@ async def sync_inference_weights( "endpoint_results": [], } + async def register_session(self, model_id="default", session_spec=None, materialize=False, request_id=None): + return { + "model_id": model_id, + "session_spec": session_spec or {}, + "materialize": materialize, + "registered": True, + } + async def register_adapter(self, model_id="default", lr=1e-5, request_id=None): return {"model_id": model_id, "lr": lr, "registered": True} diff --git a/src/xorl/server/backend/remote.py b/src/xorl/server/backend/remote.py index 7f4870b2..f1358687 100644 --- a/src/xorl/server/backend/remote.py +++ b/src/xorl/server/backend/remote.py @@ -20,6 +20,7 @@ ModelPassData, OptimStepData, RegisterAdapterData, + RegisterSessionData, SaveFullWeightsData, SaveLoraOnlyData, SaveStateData, @@ -225,7 +226,14 @@ async def _execute( 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 + self, + batches, + loss_fn="causallm_loss", + loss_fn_params=None, + model_id=None, + routed_experts=None, + routed_expert_logits=None, + request_id=None, ): return await self._execute( "forward_backward", @@ -235,11 +243,21 @@ async def forward_backward( loss_fn_params=loss_fn_params, model_id=model_id, routed_experts=routed_experts, + routed_expert_logits=routed_expert_logits, ), request_id=request_id, ) - 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, + routed_experts=None, + routed_expert_logits=None, + request_id=None, + ): return await self._execute( "forward", ModelPassData( @@ -247,6 +265,8 @@ async def forward(self, batches, loss_fn="causallm_loss", loss_fn_params=None, m loss_fn=loss_fn, loss_fn_params=loss_fn_params, model_id=model_id, + routed_experts=routed_experts, + routed_expert_logits=routed_expert_logits, ), request_id=request_id, ) @@ -335,7 +355,7 @@ async def sync_inference_weights( buffer_size_mb=1024, sync_method="nccl_broadcast", flush_cache=False, - pause_mode="retract", + pause_mode="in_place", weight_version=None, quantization=None, request_id=None, @@ -355,7 +375,19 @@ async def sync_inference_weights( quantization=quantization, ), request_id=request_id, - timeout=600.0, + timeout=self.operation_timeout, + ) + + async def register_session(self, model_id="default", session_spec=None, materialize=False, request_id=None): + return await self._execute( + "register_session", + RegisterSessionData( + model_id=model_id, + session_spec=session_spec or {}, + materialize=materialize, + ), + request_id=request_id, + timeout=60.0, ) async def register_adapter(self, model_id="default", lr=1e-5, request_id=None): diff --git a/src/xorl/server/launcher.py b/src/xorl/server/launcher.py index e56d85a1..a7dbed69 100644 --- a/src/xorl/server/launcher.py +++ b/src/xorl/server/launcher.py @@ -34,7 +34,7 @@ from contextlib import asynccontextmanager, closing from dataclasses import fields from pathlib import Path -from typing import Dict, List, Optional, Tuple +from typing import Any, Dict, List, Optional, Tuple import requests import uvicorn @@ -43,6 +43,7 @@ from xorl.server.api_server.server import APIServer from xorl.server.orchestrator.orchestrator import Orchestrator from xorl.server.server_arguments import ServerArguments +from xorl.server.session_spec import build_default_session_spec from xorl.server.utils.network import read_address_file @@ -126,15 +127,24 @@ def find_free_ports(count: int, start_port: int = 50000) -> List[int]: Returns: List of free port numbers """ + end_port = 60000 + candidates = list(range(start_port, end_port)) + random.shuffle(candidates) + ports = [] - current_start = start_port + for port in candidates: + with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as sock: + try: + sock.bind(("", port)) + sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + except OSError: + continue - for _ in range(count): - port = find_free_port(current_start) ports.append(port) - current_start = port + 1 + if len(ports) == count: + return ports - return ports + raise RuntimeError(f"Could not find {count} free ports in range {start_port}-{end_port}") # ============================================================================ @@ -245,10 +255,15 @@ def run_api_server( default_timeout: float = 120.0, output_dir: str = "outputs", base_model: Optional[str] = None, + default_session_spec: Optional[Dict[str, Any]] = None, + server_lora_config: Optional[Dict[str, Any]] = None, + max_lora_rank: Optional[int] = None, storage_limit: str = "10TB", idle_session_timeout: float = 7200.0, skip_initial_checkpoint: bool = False, sync_inference_method: str = "nccl_broadcast", + train_config: Optional[Dict[str, Any]] = None, + lora_config: Optional[Dict[str, Any]] = None, ): """ Run the API Server in a separate process. @@ -266,6 +281,8 @@ def run_api_server( idle_session_timeout: Idle session timeout in seconds. Default: 7200.0 (2 hours). 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'. + train_config: Train config defaults for per-session optimizer specs. + lora_config: LoRA config defaults and limits for per-session adapter specs. """ # Setup logging for this process @@ -306,10 +323,15 @@ async def custom_lifespan(app): default_timeout=default_timeout, output_dir=output_dir, base_model=base_model, + default_session_spec=default_session_spec, + server_lora_config=server_lora_config, + max_lora_rank=max_lora_rank, storage_limit=storage_limit, idle_session_timeout=idle_session_timeout, skip_initial_checkpoint=skip_initial_checkpoint, sync_inference_method=sync_inference_method, + train_config=train_config, + lora_config=lora_config, ) await _state_module.api_server.start() yield @@ -658,6 +680,29 @@ def __init__( self.base_model = None logger.info("No base_model configured (will not validate create_model requests)") + self.server_lora_config: Dict[str, Any] = {} + self.default_session_spec: Optional[Dict[str, Any]] = None + self.max_lora_rank: Optional[int] = None + if self.server_args: + config_dict = self.server_args.to_config_dict() + self.server_lora_config = config_dict.get("lora", {}) + self.max_lora_rank = self.server_lora_config.get( + "max_lora_rank", + self.server_lora_config.get("lora_rank"), + ) + if self.server_lora_config.get("enable_lora", False): + self.default_session_spec = build_default_session_spec( + base_model=self.base_model or self.server_args.model_path, + train_config=config_dict.get("train", {}), + lora_config=self.server_lora_config, + ) + logger.info( + "Using default multi-adapter session spec: " + f"rank={self.default_session_spec['lora_config']['lora_rank']}, " + f"alpha={self.default_session_spec['lora_config']['lora_alpha']}, " + f"optimizer={self.default_session_spec['optimizer_config']['type']}" + ) + # Storage limit - prefer from ServerArguments if available if self.server_args: self.storage_limit = self.server_args.storage_limit @@ -878,10 +923,14 @@ def _launch_workers_with_torchrun(self): stale_address_file.unlink() logger.info(f"Removed stale address file: {stale_address_file}") - # Build torchrun command — use the same Python environment as the launcher - torchrun_bin = os.path.join(os.path.dirname(sys.executable), "torchrun") + # Build torchrun command through the active Python executable. In + # mounted worktrees the torchrun console script can have a host-path + # shebang that does not exist inside the container, while + # ``python -m torch.distributed.run`` remains relocatable. cmd = [ - torchrun_bin, + sys.executable, + "-m", + "torch.distributed.run", f"--nnodes={self.nnodes}", f"--nproc-per-node={self.nproc_per_node}", f"--master-addr={self.master_addr}", @@ -1034,6 +1083,12 @@ def start(self): logger.info("✓ Engine Core process started, waiting for full initialization...") # Start API Server (connects to engine) + server_config = self.server_args.to_config_dict() if self.server_args else {} + api_train_config = server_config.get("train", {}) + api_lora_config = server_config.get("lora", {}) + skip_initial_checkpoint = self.server_args.skip_initial_checkpoint if self.server_args else False + sync_inference_method = self.server_args.sync_inference_method if self.server_args else "nccl_broadcast" + logger.info("Starting API Server...") logger.info(f" output_dir: {self.output_dir}") logger.info(f" base_model: {self.base_model}") @@ -1050,10 +1105,15 @@ def start(self): self.operation_timeout, self.output_dir, self.base_model, + self.default_session_spec, + self.server_lora_config, + self.max_lora_rank, self.storage_limit, self.idle_session_timeout, - self.server_args.skip_initial_checkpoint, - self.server_args.sync_inference_method, + skip_initial_checkpoint, + sync_inference_method, + api_train_config, + api_lora_config, ), name="APIServer", ) @@ -1078,7 +1138,7 @@ def start(self): raise RuntimeError(f"Engine Core failed during initialization (exit code: {exit_code})") # Save initial checkpoint (000) to capture the model state before any training - if not self.server_args.skip_initial_checkpoint: + if not skip_initial_checkpoint: self._save_initial_checkpoint() else: logger.info("Skipping initial checkpoint save (skip_initial_checkpoint=true)") diff --git a/src/xorl/server/orchestrator/orchestrator.py b/src/xorl/server/orchestrator/orchestrator.py index 33999b89..5d6552a3 100644 --- a/src/xorl/server/orchestrator/orchestrator.py +++ b/src/xorl/server/orchestrator/orchestrator.py @@ -30,7 +30,7 @@ │ ┌─────────────────────────────────────────────────────────────────────┐ │ │ │ API Server Communication │ │ │ │ │ │ -│ │ INPUT Socket (DEALER connect) OUTPUT Socket (PUSH bind) │ │ +│ │ INPUT Socket (DEALER connect) OUTPUT Socket (PUSH connect) │ │ │ │ ↓ ↑ │ │ │ │ input_queue (Queue) output_queue (Queue) │ │ │ │ ↓ ↑ │ │ @@ -107,7 +107,7 @@ 4. **Response:** ``` - output_queue → output_socket (PUSH) → API Server + output_queue → output_socket (PUSH connect) → API Server PULL Scheduler.mark_completed() [processing → completed] ``` @@ -227,7 +227,7 @@ def __init__( Args: input_addr: ZMQ address for input DEALER socket (API server requests) - output_addr: ZMQ address for output PUSH socket (API server responses) + output_addr: ZMQ address for output PUSH socket (connects to API server response endpoint) engine_identity: Identity for DEALER socket rank0_worker_address: ZMQ address of rank 0 worker (PAIR socket) num_workers: Number of workers for distributed execution @@ -498,7 +498,7 @@ def process_output_sockets(self): logger.info("Starting output socket thread...") self.output_channel = SyncPushChannel(self.output_addr) - self.output_channel.bind() + self.output_channel.connect() while self._running: try: @@ -657,6 +657,7 @@ def _handle_utility_request(self, request: OrchestratorRequest): "sleep": "execute_sleep", "wake_up": "execute_wake_up", "sync_inference_weights": "execute_sync_inference_weights", + "register_session": "execute_register_session", "register_adapter": "execute_register_adapter", "save_adapter_state": "execute_save_adapter_state", "load_adapter_state": "execute_load_adapter_state", diff --git a/src/xorl/server/orchestrator/packing.py b/src/xorl/server/orchestrator/packing.py index 1fa1791c..79ddf952 100644 --- a/src/xorl/server/orchestrator/packing.py +++ b/src/xorl/server/orchestrator/packing.py @@ -57,6 +57,41 @@ logger = logging.getLogger(__name__) +OPD_TOKEN_ALIGNED_FIELDS = ("teacher_ids", "teacher_cache_indices", "teacher_weights", "teacher_hidden_states") + + +def shift_opd_token_aligned_fields( + flattened_datum: Dict[str, Any], + original_seq_len: int, + shifted_seq_len: int, + sample_idx: int, +) -> None: + """Trim OPD per-token fields when HF-style causal shifting trims input_ids. + + OPD compares the student's hidden state at each retained input position with + the teacher hidden state for the same context position. Therefore fields that + index or carry teacher context activations drop the final token when + input_ids becomes input_ids[:-1]. + """ + for key in OPD_TOKEN_ALIGNED_FIELDS: + if key not in flattened_datum: + continue + value = flattened_datum[key] + if hasattr(value, "tolist"): + value = value.tolist() + if not isinstance(value, list): + continue + if len(value) == shifted_seq_len: + flattened_datum[key] = value + continue + if len(value) != original_seq_len: + raise ValueError( + f"Sample {sample_idx}: {key} length ({len(value)}) must match either shifted length " + f"({shifted_seq_len}) or original length ({original_seq_len})" + ) + flattened_datum[key] = value[:-1] + + def apply_weights_to_labels( labels: List[int], weights: Optional[List[float]], @@ -580,6 +615,12 @@ def _add_sample_to_packed_batch( f"Sample {sample_idx}: Applying token shifting (HF format detected). " f"Original len={len(input_ids)}, shifted len={len(input_ids) - 1}" ) + shift_opd_token_aligned_fields( + flattened_datum, + original_seq_len=len(input_ids), + shifted_seq_len=len(input_ids) - 1, + sample_idx=sample_idx, + ) input_ids = input_ids[:-1] labels = labels[1:] if weights is not None: @@ -618,9 +659,19 @@ def _add_sample_to_packed_batch( batch["labels"].extend(labels) + # OPD convenience fields: allow sample-level teacher metadata and expand + # it to token-aligned sequence fields before generic field preservation. + teacher_id = flattened_datum.get("teacher_id") + if teacher_id is not None and "teacher_ids" not in flattened_datum: + batch.setdefault("teacher_ids", []).extend([int(teacher_id)] * seq_len) + + teacher_weight = flattened_datum.get("teacher_weight") + if teacher_weight is not None and "teacher_weights" not in flattened_datum: + batch.setdefault("teacher_weights", []).extend([float(teacher_weight)] * seq_len) + # Handle other sequence fields (logprobs, advantages, etc.) for key, value in flattened_datum.items(): - if key not in ["input_ids", "position_ids", "labels", "weights"]: + if key not in ["input_ids", "position_ids", "labels", "weights", "teacher_id", "teacher_weight"]: if key not in batch: batch[key] = [] if hasattr(value, "tolist"): @@ -688,7 +739,13 @@ def _finalize_packed_batch(self, batch: Dict[str, Any]) -> None: continue if isinstance(value, list) and len(value) == 1 and isinstance(value[0], list): if len(value[0]) == seq_len: - value[0].extend([0] * pad_length) + if key == "teacher_hidden_states": + if not value[0] or not isinstance(value[0][0], list): + raise ValueError("teacher_hidden_states must be a sequence of hidden vectors") + hidden_dim = len(value[0][0]) + value[0].extend([[0.0] * hidden_dim for _ in range(pad_length)]) + else: + value[0].extend([0] * pad_length) # For non-SP cases, pre-compute Flash Attention kwargs from position_ids. # (SP cases handle this in TextSequenceShardCollator after SP padding.) @@ -716,7 +773,14 @@ def _add_sample_to_batch( datum: Dict[str, Any], sample_idx: int, ) -> None: - """Add a sample to a micro-batch.""" + """Add a sample to a non-packed micro-batch. + + This mirrors the packed path's token-shifting semantics: + - HF-format datums (`input_ids` + full-length `labels`) are shifted to + next-token prediction with `input_ids[:-1]` and `labels[1:]`. + - xorl_client/RL datums (`input_ids` + `target_tokens`) are already + shifted and are used as-is. + """ # Handle nested datum structure: flatten model_input and loss_fn_inputs flattened_datum = {} if "model_input" in datum: @@ -739,19 +803,13 @@ def _add_sample_to_batch( if not isinstance(input_ids, list): 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) - # Extract or generate position_ids 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) else: - # Auto-generate position_ids: [0, 1, 2, ..., seq_len-1] - position_ids = list(range(seq_len)) - - batch["position_ids"].append(position_ids) + position_ids = list(range(len(input_ids))) # Extract labels if present (or use target_tokens for RL) if "labels" in flattened_datum: @@ -767,31 +825,66 @@ def _add_sample_to_batch( # No labels for this sample labels = [] + weights = flattened_datum.get("weights") + if weights is not None and not isinstance(weights, list): + weights = weights.tolist() if hasattr(weights, "tolist") else list(weights) + + advantages = flattened_datum.get("advantages") + if advantages is not None and not isinstance(advantages, list): + advantages = advantages.tolist() if hasattr(advantages, "tolist") else list(advantages) + + # Detect if tokens are already shifted (xorl_client API format) + is_already_shifted = "target_tokens" in flattened_datum and len(input_ids) == len(labels) + if labels and not is_already_shifted and len(input_ids) == len(labels): + logger.warning( + "Sample %s has labels with the same length as input_ids; treating it as HF-format data " + "and shifting for next-token prediction. Use target_tokens for already shifted targets.", + sample_idx, + ) + input_ids = input_ids[:-1] + position_ids = position_ids[:-1] + labels = labels[1:] + if weights is not None: + weights = weights[1:] + if advantages is not None: + advantages = advantages[1:] + + if advantages is not None: + flattened_datum["advantages"] = advantages + + batch["input_ids"].append(input_ids) + batch["position_ids"].append(position_ids) + # Apply weights mask to labels if weights field is present # weights=0 -> labels=-100 (IGNORE_INDEX), weights=1 -> labels unchanged if labels: # Only apply if we have labels - 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) labels = apply_weights_to_labels(labels, weights, sample_idx) # Apply advantages mask to labels if advantages field is present # For RL losses, advantages=0 indicates prompt tokens where we don't compute loss # advantages=0 -> labels=-100 (IGNORE_INDEX) - 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) labels = apply_advantages_to_labels(labels, advantages, sample_idx) batch["labels"].append(labels) + # OPD convenience fields: expand sample-level metadata to token-aligned + # sequences so downstream tensor conversion/sharding can treat them like + # labels and logprobs. + teacher_id = flattened_datum.get("teacher_id") + if teacher_id is not None and "teacher_ids" not in flattened_datum: + batch.setdefault("teacher_ids", []).append([int(teacher_id)] * seq_len) + + teacher_weight = flattened_datum.get("teacher_weight") + if teacher_weight is not None and "teacher_weights" not in flattened_datum: + batch.setdefault("teacher_weights", []).append([float(teacher_weight)] * seq_len) + # Preserve all other fields from loss_fn_inputs (logprobs, advantages, target_tokens, etc.) # Note: Keep target_tokens separate even if we used it as labels # Note: Exclude weights as it has been applied to labels for key, value in flattened_datum.items(): - if key not in ["input_ids", "position_ids", "labels", "weights"]: + if key not in ["input_ids", "position_ids", "labels", "weights", "teacher_id", "teacher_weight"]: # Initialize list for this field if not present if key not in batch: batch[key] = [] diff --git a/src/xorl/server/orchestrator/request_processor.py b/src/xorl/server/orchestrator/request_processor.py index 60de2448..e89f6a0d 100644 --- a/src/xorl/server/orchestrator/request_processor.py +++ b/src/xorl/server/orchestrator/request_processor.py @@ -55,6 +55,7 @@ ModelPassData, OptimStepData, RegisterAdapterData, + RegisterSessionData, SaveFullWeightsData, SaveLoraOnlyData, SaveStateData, @@ -163,10 +164,20 @@ async def _execute_model_pass( data = p.data loss_fn = p.loss_fn loss_fn_params = p.loss_fn_params or {} + routed_experts = p.routed_experts + routed_expert_logits = p.routed_expert_logits if not data: raise ValueError("data or datum_list must be provided") + if loss_fn == "opd_loss" and loss_fn_params.get("opd_sort_by_teacher", True): + order = sorted(range(len(data)), key=lambda i: self._teacher_sort_key(data[i])) + data = [data[i] for i in order] + if routed_experts is not None: + routed_experts = [routed_experts[i] for i in order] + if routed_expert_logits is not None: + routed_expert_logits = [routed_expert_logits[i] for i in order] + # Pack samples into batches logger.debug(f"Packing {len(data)} datum into batches for {op_name} request {request.request_id}") batches = pack_samples( @@ -195,10 +206,10 @@ async def _execute_model_pass( loss_fn=loss_fn, loss_fn_params=loss_fn_params, model_id=p.model_id, + routed_experts=routed_experts, + routed_expert_logits=routed_expert_logits, request_id=request.request_id, ) - if op_name == "forward_backward": - kwargs["routed_experts"] = p.routed_experts result = await backend_method(**kwargs) @@ -212,12 +223,14 @@ async def _execute_model_pass( "loss": loss, "valid_tokens": tokens, "success": True, - "execution_time": result.get("execution_time", 0.0), + "execution_time": result.get( + "execution_time", result.get("forward_backward_time", result.get("forward_time", 0.0)) + ), } - # Add IS metrics (KL divergence, ratio stats, etc.) + # Add loss-specific metrics (IS/KL divergence, OPD KL stats, ratio stats, etc.) for key in result: - if key.startswith("is_"): + if key.startswith(("is_", "opd_")): output_dict[key] = result[key] # Pass through expert load summary for MoE models @@ -229,6 +242,17 @@ async def _execute_model_pass( output_dict["auto_loaded"] = True output_dict["auto_load_path"] = result.get("auto_load_path") + # Forward-only teacher prefill path writes activation caches to + # shared storage and returns the cache metadata here. + for key in ( + "teacher_hidden_cache", + "teacher_prefill_tokens", + "teacher_prefill_forward_compute_s", + "teacher_hidden_cache_write_s", + ): + if key in result: + output_dict[key] = result[key] + # Unpack per-token outputs if present (tinker API compatibility) if "packed_logprobs" in result and "packed_position_ids" in result: output_dict["per_sample_outputs"] = self._unpack_per_sample_outputs(result, batches) @@ -267,6 +291,28 @@ async def _execute_model_pass( error=error_msg, ) + @staticmethod + def _teacher_sort_key(datum: Dict[str, Any]) -> int: + flattened: Dict[str, Any] = {} + if isinstance(datum.get("model_input"), dict): + flattened.update(datum["model_input"]) + if isinstance(datum.get("loss_fn_inputs"), dict): + flattened.update(datum["loss_fn_inputs"]) + for key, value in datum.items(): + if key not in ("model_input", "loss_fn_inputs"): + flattened[key] = value + + teacher_id = flattened.get("teacher_id") + if teacher_id is None: + teacher_ids = flattened.get("teacher_ids") + if hasattr(teacher_ids, "reshape"): + teacher_ids = teacher_ids.reshape(-1).tolist() + if isinstance(teacher_ids, list) and teacher_ids: + while isinstance(teacher_ids[0], list) and teacher_ids[0]: + teacher_ids = teacher_ids[0] + teacher_id = teacher_ids[0] + return int(teacher_id) if teacher_id is not None else 0 + @staticmethod def _unpack_per_sample_outputs(result: Dict, batches: list) -> list: """Unpack packed per-token outputs into per-sample lists. @@ -571,6 +617,15 @@ async def execute_sync_inference_weights(self, request: OrchestratorRequest) -> if not p.endpoints: raise ValueError("inference endpoints must be provided") + group_name = p.group_name + if p.sync_method == "nccl_broadcast": + # NCCL weight sync uses abort-based teardown to avoid cooperative + # shutdown hangs. PyTorch can keep the group name reserved after + # abort(), so scope NCCL group names to the request. + request_token = "".join(ch if ch.isalnum() or ch == "_" else "_" for ch in str(request.request_id)) + request_token = request_token.strip("_")[:32] or f"{time.time_ns():x}" + group_name = f"{p.group_name}_{request_token}" + def build_output(result): return [ { @@ -580,6 +635,8 @@ def build_output(result): "total_bytes": result.get("total_bytes", 0), "num_parameters": result.get("num_parameters", 0), "num_buckets": result.get("num_buckets", 0), + "timing_breakdown": result.get("timing_breakdown", {}), + "p2p_rank_summaries": result.get("p2p_rank_summaries", []), "endpoint_results": result.get("endpoint_results", []), "execution_time": result.get("execution_time", 0.0), } @@ -592,7 +649,7 @@ def build_output(result): endpoints=p.endpoints, master_address=p.master_address, master_port=p.master_port, - group_name=p.group_name, + group_name=group_name, buffer_size_mb=p.buffer_size_mb, sync_method=p.sync_method, flush_cache=p.flush_cache, @@ -624,6 +681,26 @@ def build_output(result): build_output, ) + async def execute_register_session(self, request: OrchestratorRequest) -> OrchestratorOutputs: + """Register a normalized session runtime spec on workers.""" + p: RegisterSessionData = request.payload + + def build_output(result): + return {"result": result} + + return await self._execute_operation( + request, + "register_session", + self.backend.register_session( + model_id=p.model_id, + session_spec=p.session_spec, + materialize=p.materialize, + request_id=request.request_id, + ), + OutputType.REGISTER_SESSION, + build_output, + ) + async def execute_save_adapter_state(self, request: OrchestratorRequest) -> OrchestratorOutputs: """Execute save adapter state on workers.""" p: AdapterStateData = request.payload diff --git a/src/xorl/server/protocol/__init__.py b/src/xorl/server/protocol/__init__.py index ce3b9f45..48ee61f0 100644 --- a/src/xorl/server/protocol/__init__.py +++ b/src/xorl/server/protocol/__init__.py @@ -43,6 +43,7 @@ OperationPayload, OptimStepData, RegisterAdapterData, + RegisterSessionData, SaveFullWeightsData, SaveLoraOnlyData, SaveStateData, diff --git a/src/xorl/server/protocol/api_orchestrator.py b/src/xorl/server/protocol/api_orchestrator.py index 310b4921..67c5580e 100644 --- a/src/xorl/server/protocol/api_orchestrator.py +++ b/src/xorl/server/protocol/api_orchestrator.py @@ -63,6 +63,7 @@ class OutputType(str, Enum): WAKE_UP = "wake_up" HEALTH_CHECK = "health_check" SYNC_INFERENCE_WEIGHTS = "sync_inference_weights" + REGISTER_SESSION = "register_session" REGISTER_ADAPTER = "register_adapter" SAVE_ADAPTER_STATE = "save_adapter_state" LOAD_ADAPTER_STATE = "load_adapter_state" @@ -430,6 +431,8 @@ def create_sync_weights_output( total_bytes: int = 0, num_parameters: int = 0, num_buckets: int = 0, + timing_breakdown: Optional[Dict[str, float]] = None, + p2p_rank_summaries: Optional[List[Dict[str, Any]]] = None, endpoint_results: Optional[List[Dict[str, Any]]] = None, error: Optional[str] = None, ) -> OrchestratorOutputs: @@ -444,6 +447,8 @@ def create_sync_weights_output( total_bytes=total_bytes, num_parameters=num_parameters, num_buckets=num_buckets, + timing_breakdown=timing_breakdown or {}, + p2p_rank_summaries=p2p_rank_summaries or [], endpoint_results=endpoint_results or [], ) diff --git a/src/xorl/server/protocol/operations.py b/src/xorl/server/protocol/operations.py index 395ecca2..5739b248 100644 --- a/src/xorl/server/protocol/operations.py +++ b/src/xorl/server/protocol/operations.py @@ -104,7 +104,7 @@ class SyncWeightsData: buffer_size_mb: int = 1024 sync_method: str = "nccl_broadcast" flush_cache: bool = False - pause_mode: str = "retract" + pause_mode: str = "in_place" weight_version: Optional[str] = None quantization: Optional[Dict[str, Any]] = None @@ -117,6 +117,15 @@ class RegisterAdapterData: lr: float = 1e-5 +@dataclass +class RegisterSessionData: + """Payload for register_session operations.""" + + model_id: str = "default" + session_spec: Dict[str, Any] = field(default_factory=dict) + materialize: bool = False + + @dataclass class AdapterStateData: """Payload for save_adapter_state / load_adapter_state operations.""" @@ -164,6 +173,7 @@ class EmptyData: SaveFullWeightsData, SyncWeightsData, RegisterAdapterData, + RegisterSessionData, AdapterStateData, KillSessionData, AbortData, @@ -182,6 +192,7 @@ class EmptyData: "save_full_weights": SaveFullWeightsData, "sync_inference_weights": SyncWeightsData, "register_adapter": RegisterAdapterData, + "register_session": RegisterSessionData, "save_adapter_state": AdapterStateData, "load_adapter_state": AdapterStateData, "kill_session": KillSessionData, diff --git a/src/xorl/server/protocol/orchestrator_runner.py b/src/xorl/server/protocol/orchestrator_runner.py index 640d79f7..2be40c90 100644 --- a/src/xorl/server/protocol/orchestrator_runner.py +++ b/src/xorl/server/protocol/orchestrator_runner.py @@ -62,6 +62,7 @@ class MessageType(str, Enum): WAKE_UP = "wake_up" HEALTH_CHECK = "health_check" SYNC_INFERENCE_WEIGHTS = "sync_inference_weights" + REGISTER_SESSION = "register_session" REGISTER_ADAPTER = "register_adapter" SAVE_ADAPTER_STATE = "save_adapter_state" LOAD_ADAPTER_STATE = "load_adapter_state" diff --git a/src/xorl/server/runner/adapters/adapter_coordinator.py b/src/xorl/server/runner/adapters/adapter_coordinator.py index d9230653..4f4678cd 100644 --- a/src/xorl/server/runner/adapters/adapter_coordinator.py +++ b/src/xorl/server/runner/adapters/adapter_coordinator.py @@ -9,22 +9,37 @@ live here. The RunnerDispatcher delegates to this class. """ +from __future__ import annotations + +import json import logging import os -from typing import Any, Dict, Optional, Tuple +import time +from copy import deepcopy +from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple +import torch import torch.distributed as dist +from safetensors.torch import load_file as safetensors_load_file +from xorl.lora.utils import convert_peft_lora_state_dict, get_lora_tensor_shard_specs from xorl.server.protocol.operations import ( AdapterStateData, KillSessionData, RegisterAdapterData, + RegisterSessionData, ) -from xorl.server.runner.model_runner import ModelRunner +from xorl.server.session_spec import load_session_spec_from_checkpoint + + +if TYPE_CHECKING: + from xorl.server.runner.model_runner import ModelRunner logger = logging.getLogger(__name__) +_ADAPTER_STATE_LOAD_MODES = {"all_ranks", "rank0_broadcast"} + class AdapterCoordinator: """Coordinates multi-rank LoRA adapter operations. @@ -46,6 +61,15 @@ def __init__( self.world_size = world_size self.cpu_group = cpu_group + def _validate_pipeline_parallel_broadcast_safe(self) -> None: + """Reject pipeline-parallel topologies for broadcast-based adapter coordination.""" + pipeline_parallel_size = int(getattr(self.trainer, "train_config", {}).get("pipeline_parallel_size", 1)) + if pipeline_parallel_size > 1 and self.world_size > 1: + raise RuntimeError( + "pipeline_parallel_size > 1 is not supported with multi-adapter LoRA server training. " + "Adapter coordination currently assumes identical local LoRA layouts on every rank." + ) + # ======================================================================== # Adapter Broadcast # ======================================================================== @@ -60,6 +84,7 @@ def broadcast_adapter_state(self, model_id: str, default_lr: float) -> None: """ if self.world_size <= 1: return + self._validate_pipeline_parallel_broadcast_safe() adapter_state = self.trainer.adapter_manager.get_adapter_state(model_id) @@ -83,10 +108,436 @@ def broadcast_adapter_state(self, model_id: str, default_lr: float) -> None: if self.rank != 0 and metadata[0]: adapter_state.global_step = metadata[0].get("global_step", 0) adapter_state.global_forward_backward_step = metadata[0].get("global_forward_backward_step", 0) - adapter_state.lr = metadata[0].get("lr", default_lr) + self.trainer.adapter_manager.set_lr(model_id, metadata[0].get("lr", default_lr)) + adapter_state.last_access_time = time.time() logger.debug(f"Rank {self.rank}: Broadcast adapter state for model_id={model_id}") + @staticmethod + def _optimizer_state_to_cpu(value: Any) -> Any: + """Recursively move optimizer state dict tensors to CPU for object broadcast.""" + if isinstance(value, torch.Tensor): + return value.detach().cpu() + if isinstance(value, dict): + return {k: AdapterCoordinator._optimizer_state_to_cpu(v) for k, v in value.items()} + if isinstance(value, list): + return [AdapterCoordinator._optimizer_state_to_cpu(v) for v in value] + if isinstance(value, tuple): + return tuple(AdapterCoordinator._optimizer_state_to_cpu(v) for v in value) + return value + + def broadcast_adapter_optimizer_state(self, model_id: str) -> None: + """Broadcast adapter optimizer state from rank 0 to all other ranks.""" + if self.world_size <= 1: + return + self._validate_pipeline_parallel_broadcast_safe() + + adapter_state = self.trainer.adapter_manager.get_adapter_state(model_id) + optimizer_state = [None] + if self.rank == 0: + optimizer_state[0] = self._optimizer_state_to_cpu(adapter_state.optimizer.state_dict()) + + dist.broadcast_object_list(optimizer_state, src=0, group=self.cpu_group) + + if self.rank != 0 and optimizer_state[0] is not None: + adapter_state.optimizer.load_state_dict(optimizer_state[0]) + + def _has_ep_sharded_adapter_params(self, model_id: str) -> bool: + """Return whether the resident adapter has LoRA tensors sharded by EP.""" + adapter_manager = self.trainer.adapter_manager + model = getattr(adapter_manager, "model", getattr(self.trainer, "model", None)) + if adapter_manager is None or model is None or not adapter_manager.has_adapter(model_id): + return False + + canonical_name = getattr(adapter_manager, "_canonical_lora_param_name", lambda name: name) + state = adapter_manager.get_adapter_state(model_id) + requested_names = {canonical_name(name) for name in state.lora_params} + return bool(get_lora_tensor_shard_specs(model, names=requested_names)) + + def _expected_adapter_param_maps(self, model_id: str) -> Tuple[Dict[str, str], Dict[str, torch.Size]]: + """Build canonical-name maps for the live adapter tensors.""" + adapter_manager = self.trainer.adapter_manager + state = adapter_manager.get_adapter_state(model_id) + expected_param_map: Dict[str, str] = {} + expected_shapes: Dict[str, torch.Size] = {} + + for actual_name, param in state.lora_params.items(): + canonical_name = adapter_manager._canonical_lora_param_name(actual_name) + if canonical_name in expected_param_map and expected_param_map[canonical_name] != actual_name: + raise ValueError( + f"Live adapter contains duplicate LoRA tensors after canonicalization. param={canonical_name!r}" + ) + expected_param_map[canonical_name] = actual_name + expected_shapes[canonical_name] = param.shape + + return expected_param_map, expected_shapes + + @staticmethod + def _strip_optimizer_config(session_spec: Dict[str, Any]) -> Dict[str, Any]: + stripped = deepcopy(session_spec) + stripped.pop("optimizer_config", None) + return stripped + + @staticmethod + def _strip_optimizer_learning_rate(session_spec: Dict[str, Any]) -> Dict[str, Any]: + stripped = deepcopy(session_spec) + optimizer_config = stripped.get("optimizer_config") + if isinstance(optimizer_config, dict): + optimizer_config.pop("learning_rate", None) + return stripped + + def _validate_broadcast_checkpoint_session_spec( + self, + model_id: str, + checkpoint_session_spec: Optional[Dict[str, Any]], + *, + load_optimizer: bool, + lr: Optional[float], + ) -> None: + """Validate rank-0 broadcast checkpoint metadata before applying tensors.""" + if not isinstance(checkpoint_session_spec, dict): + raise ValueError("Rank-0 broadcast adapter payload did not include a valid checkpoint session spec") + + registered_session_spec = self.trainer.get_lora_session_spec(model_id) + checkpoint_spec_for_compare = checkpoint_session_spec + registered_spec_for_compare = registered_session_spec + if lr is not None: + checkpoint_spec_for_compare = self._strip_optimizer_learning_rate(checkpoint_spec_for_compare) + registered_spec_for_compare = self._strip_optimizer_learning_rate(registered_spec_for_compare) + + if load_optimizer: + specs_match = checkpoint_spec_for_compare == registered_spec_for_compare + mismatch_context = "registered multi-adapter session" + else: + specs_match = self._strip_optimizer_config(checkpoint_spec_for_compare) == self._strip_optimizer_config( + registered_spec_for_compare + ) + mismatch_context = "registered multi-adapter session for weights-only restore" + + if not specs_match: + raise ValueError( + "Checkpoint session spec does not match the " + f"{mismatch_context}. checkpoint={checkpoint_session_spec!r}, " + f"current={registered_session_spec!r}" + ) + + def _rank0_load_adapter_checkpoint_payload(self, model_id: str, path: str, load_optimizer: bool) -> Dict[str, Any]: + """Load adapter checkpoint tensors on rank 0 and broadcast them as a CPU payload.""" + payload = [None] + + if self.rank == 0: + try: + if not os.path.exists(path): + raise FileNotFoundError(f"Checkpoint path does not exist: {path}") + + validate_checkpoint = getattr(self.trainer.adapter_manager, "_validate_checkpoint_adapter_config", None) + if validate_checkpoint is not None: + validate_checkpoint(path) + + registered_session_spec = self.trainer.get_lora_session_spec(model_id) + checkpoint_session_spec = load_session_spec_from_checkpoint( + path, + fallback_base_model=registered_session_spec.get("base_model"), + fallback_session_spec=registered_session_spec, + ) + + metadata_path = os.path.join(path, "metadata.json") + if os.path.exists(metadata_path): + with open(metadata_path, "r") as f: + metadata = json.load(f) + else: + metadata = {} + + weights_path = os.path.join(path, "adapter_model.safetensors") + if not os.path.exists(weights_path): + raise FileNotFoundError(f"Weights file not found: {weights_path}") + + loaded_weights = safetensors_load_file(weights_path) + payload[0] = { + "error": None, + "session_spec": checkpoint_session_spec, + "metadata": metadata, + "weights": {name: tensor.cpu() for name, tensor in loaded_weights.items()}, + "optimizer_present": load_optimizer and os.path.exists(os.path.join(path, "optimizer.pt")), + } + except Exception as e: + payload[0] = {"error": str(e)} + + dist.broadcast_object_list(payload, src=0, group=self.cpu_group) + if payload[0].get("error"): + raise RuntimeError(payload[0]["error"]) + return payload[0] + + def _apply_broadcast_adapter_checkpoint_payload( + self, + model_id: str, + payload: Dict[str, Any], + *, + load_optimizer: bool, + lr: Optional[float], + ) -> None: + """Convert a rank0-broadcast checkpoint payload into this rank's local adapter tensors.""" + adapter_manager = self.trainer.adapter_manager + state = adapter_manager.get_adapter_state(model_id) + self._validate_broadcast_checkpoint_session_spec( + model_id, + payload.get("session_spec"), + load_optimizer=load_optimizer, + lr=lr, + ) + expected_param_map, expected_shapes = self._expected_adapter_param_maps(model_id) + expected_shard_specs = get_lora_tensor_shard_specs(adapter_manager.model, names=expected_shapes.keys()) + + converted_weights = convert_peft_lora_state_dict( + payload["weights"], + expected_shapes=expected_shapes, + expected_shard_specs=expected_shard_specs, + ) + + checkpoint_tensors: Dict[str, torch.Tensor] = {} + for converted_name, weight in converted_weights.items(): + canonical_name = adapter_manager._canonical_lora_param_name(converted_name) + if canonical_name in checkpoint_tensors: + raise ValueError( + f"Checkpoint contains duplicate LoRA tensors after canonicalization. param={canonical_name!r}" + ) + checkpoint_tensors[canonical_name] = weight + + expected_param_names = set(expected_param_map) + checkpoint_param_names = set(checkpoint_tensors) + missing_param_names = sorted(expected_param_names - checkpoint_param_names) + unexpected_param_names = sorted(checkpoint_param_names - expected_param_names) + if missing_param_names or unexpected_param_names: + raise ValueError( + "Checkpoint LoRA parameter set does not match the live adapter structure. " + f"missing={missing_param_names!r}, unexpected={unexpected_param_names!r}" + ) + + for internal_name, tensor in checkpoint_tensors.items(): + target_param = state.lora_params[expected_param_map[internal_name]] + if tuple(tensor.shape) != tuple(target_param.shape): + raise ValueError( + "Checkpoint tensor shape does not match the live adapter shape. " + f"param={internal_name!r}, checkpoint={tuple(tensor.shape)!r}, " + f"live={tuple(target_param.shape)!r}" + ) + + for internal_name, tensor in checkpoint_tensors.items(): + target_param = state.lora_params[expected_param_map[internal_name]] + target_param.data.copy_(tensor.to(device=target_param.device, dtype=target_param.dtype)) + + metadata = payload.get("metadata", {}) + state.global_step = metadata.get("global_step", 0) + state.global_forward_backward_step = metadata.get("global_forward_backward_step", 0) + if lr is not None: + adapter_manager.set_lr(model_id, lr) + elif "lr" in metadata: + adapter_manager.set_lr(model_id, metadata["lr"]) + state.last_access_time = time.time() + + def _restore_ep_sharded_rank0_broadcast_adapter_state( + self, + model_id: str, + path: str, + *, + load_optimizer: bool, + lr: Optional[float], + ) -> Dict[str, Any]: + """Restore an EP-sharded LoRA adapter without broadcasting rank 0's local expert slice.""" + start_time = time.time() + payload = self._rank0_load_adapter_checkpoint_payload(model_id, path, load_optimizer) + self._apply_broadcast_adapter_checkpoint_payload( + model_id, + payload, + load_optimizer=load_optimizer, + lr=lr, + ) + + if payload.get("optimizer_present") and self.rank == 0: + logger.warning( + "Skipping optimizer restore for EP-sharded rank0_broadcast adapter load because optimizer.pt " + "contains rank-local optimizer tensors. Adapter weights and metadata were restored safely." + ) + + state = self.trainer.adapter_manager.get_adapter_state(model_id) + return { + "path": path, + "model_id": model_id, + "step": state.global_step, + "load_time": time.time() - start_time, + "success": True, + } + + def _get_adapter_state_load_mode(self) -> str: + """Return the configured adapter-state restore mode.""" + mode = getattr(self.trainer, "lora_config", {}).get("adapter_state_load_mode", "all_ranks") + if mode not in _ADAPTER_STATE_LOAD_MODES: + raise ValueError( + f"Unsupported adapter_state_load_mode: {mode!r}. " + f"Supported: {', '.join(sorted(_ADAPTER_STATE_LOAD_MODES))}." + ) + return mode + + def _sync_collective_error(self, local_error: Optional[str]) -> Optional[str]: + """Synchronize restore/registration failures before collective broadcast.""" + if self.world_size <= 1 or not dist.is_available() or not dist.is_initialized(): + return local_error + + group = self.cpu_group + if group is not None: + backend = dist.get_backend(group) + else: + backend = dist.get_backend() + device = ( + torch.device(f"cuda:{torch.cuda.current_device()}") + if backend == "nccl" and torch.cuda.is_available() + else torch.device("cpu") + ) + + has_error = torch.tensor([1 if local_error else 0], dtype=torch.int64, device=device) + dist.all_reduce(has_error, op=dist.ReduceOp.MAX, group=group) + + if has_error.item() == 0: + return None + + error_strings = [None] * self.world_size + dist.all_gather_object(error_strings, local_error or "", group=group) + errors = {i: msg for i, msg in enumerate(error_strings) if msg} + if errors: + return "; ".join(f"rank {i}: {msg}" for i, msg in errors.items()) + return local_error + + def _rollback_created_adapter(self, model_id: str, created_adapter: bool) -> None: + """Remove a newly materialized adapter after a failed restore attempt.""" + if not created_adapter or self.trainer.adapter_manager is None: + return + if not self.trainer.adapter_manager.has_adapter(model_id): + return + try: + self.trainer.adapter_manager.remove_adapter(model_id) + except Exception as e: + logger.warning(f"Rank {self.rank}: Failed to roll back adapter '{model_id}' after restore error: {e}") + + def _rollback_session_registration(self, model_id: str, *, had_session_spec: bool, had_adapter: bool) -> None: + """Remove newly installed adapter/session state after a failed registration.""" + if ( + not had_adapter + and self.trainer.adapter_manager is not None + and self.trainer.adapter_manager.has_adapter(model_id) + ): + try: + self.trainer.adapter_manager.remove_adapter(model_id) + except Exception as e: + logger.warning( + f"Rank {self.rank}: Failed to roll back adapter '{model_id}' after registration error: {e}" + ) + + if not had_session_spec and model_id in self.trainer.lora_session_specs: + self.trainer.lora_session_specs.pop(model_id, None) + + def _ensure_adapter_materialized_for_restore(self, model_id: str, lr: float) -> bool: + """Materialize a nonresident adapter and fail collectively if any rank cannot.""" + created_adapter = False + local_error = None + if not self.trainer.adapter_manager.has_adapter(model_id): + try: + self.trainer.register_lora_adapter(model_id, lr) + created_adapter = True + except Exception as e: + local_error = f"Failed to register adapter for restore: {e}" + + synced_error = self._sync_collective_error(local_error) + if synced_error: + self._rollback_created_adapter(model_id, created_adapter) + raise RuntimeError(synced_error) + + return created_adapter + + def _ensure_fresh_adapter_materialized(self, model_id: str) -> float: + """Materialize a fresh nonresident adapter and sync failures before broadcast.""" + created_adapter = False + local_error = None + default_lr = None + try: + session_spec = self.trainer.get_lora_session_spec(model_id) + default_lr = session_spec["optimizer_config"]["learning_rate"] + if not self.trainer.adapter_manager.has_adapter(model_id): + self.trainer.register_lora_adapter(model_id, default_lr) + created_adapter = True + except Exception as e: + local_error = f"Failed to register fresh adapter for model_id={model_id}: {e}" + + synced_error = self._sync_collective_error(local_error) + if synced_error: + self._rollback_created_adapter(model_id, created_adapter) + raise RuntimeError(synced_error) + + if default_lr is None: + default_lr = self.trainer.adapter_manager.get_adapter_state(model_id).lr + return default_lr + + def _restore_adapter_state( + self, + model_id: str, + path: str, + *, + load_optimizer: bool, + lr: Optional[float], + default_lr: float, + created_adapter: bool = False, + ) -> Dict[str, Any]: + """Restore adapter state using the configured rank loading strategy.""" + mode = self._get_adapter_state_load_mode() + result = None + local_error = None + + try: + if mode == "rank0_broadcast" and self.world_size > 1 and self._has_ep_sharded_adapter_params(model_id): + result = self._restore_ep_sharded_rank0_broadcast_adapter_state( + model_id=model_id, + path=path, + load_optimizer=load_optimizer, + lr=lr, + ) + elif mode == "all_ranks" or self.world_size <= 1: + result = self.trainer.load_adapter_state( + model_id=model_id, + path=path, + load_optimizer=load_optimizer, + lr=lr, + ) + elif self.rank == 0: + result = self.trainer.load_adapter_state( + model_id=model_id, + path=path, + load_optimizer=load_optimizer, + lr=lr, + ) + except Exception as e: + local_error = f"Adapter state restore failed for model_id={model_id}: {e}" + + synced_error = self._sync_collective_error(local_error) + if synced_error: + self._rollback_created_adapter(model_id, created_adapter) + raise RuntimeError(synced_error) + + # In all_ranks mode every rank has loaded its own local tensor contents. + # The EP-sharded rank0_broadcast path also materializes each rank's local + # expert slice directly from the full checkpoint tensors. + if mode == "rank0_broadcast" and not self._has_ep_sharded_adapter_params(model_id): + self.broadcast_adapter_state(model_id, default_lr) + if load_optimizer and self.world_size > 1: + self.broadcast_adapter_optimizer_state(model_id) + + if result is None: + adapter_state = self.trainer.adapter_manager.get_adapter_state(model_id) + result = { + "success": True, + "model_id": model_id, + "step": adapter_state.global_step, + } + return result + # ======================================================================== # Auto-Load Evicted Adapters # ======================================================================== @@ -109,6 +560,20 @@ def _find_evicted_checkpoint(self, model_id: str) -> Optional[str]: return evicted_path return None + def _resolve_evicted_checkpoint(self, model_id: str) -> Optional[str]: + """Resolve the evicted checkpoint path using the configured load mode.""" + if ( + self.world_size > 1 + and self.cpu_group is not None + and self._get_adapter_state_load_mode() == "rank0_broadcast" + ): + checkpoint_ref = [None] + if self.rank == 0: + checkpoint_ref[0] = self._find_evicted_checkpoint(model_id) + dist.broadcast_object_list(checkpoint_ref, src=0, group=self.cpu_group) + return checkpoint_ref[0] + return self._find_evicted_checkpoint(model_id) + def _register_fresh_adapter(self, model_id: str, lr: float = 1e-5) -> None: """Register a new adapter with fresh weights. Raises on failure. @@ -131,7 +596,12 @@ def _register_fresh_adapter(self, model_id: str, lr: float = 1e-5) -> None: f"Call /api/v1/register_adapter first." ) - def auto_load_if_evicted(self, model_id: str) -> Tuple[bool, Optional[str]]: + def auto_load_if_evicted( + self, + model_id: str, + *, + allow_fresh_materialization: bool = True, + ) -> Tuple[bool, Optional[str]]: """ Check if an adapter was evicted and auto-load from checkpoint if available. @@ -142,6 +612,8 @@ def auto_load_if_evicted(self, model_id: str) -> Tuple[bool, Optional[str]]: Args: model_id: The adapter/session ID to check + allow_fresh_materialization: When False, require a real evicted checkpoint + for nonresident adapters instead of creating fresh step-0 state. Returns: Tuple of (was_auto_loaded, checkpoint_path) @@ -155,29 +627,50 @@ def auto_load_if_evicted(self, model_id: str) -> Tuple[bool, Optional[str]]: return False, None # Look for evicted checkpoint - checkpoint_path = self._find_evicted_checkpoint(model_id) + checkpoint_path = self._resolve_evicted_checkpoint(model_id) if checkpoint_path is None: + if not allow_fresh_materialization: + raise FileNotFoundError( + f"Adapter '{model_id}' is not resident and no evicted checkpoint was found under " + f"{os.path.join(self.trainer.adapter_manager.checkpoint_dir, 'evicted', model_id)}. " + "Refusing to recreate fresh state for this operation because that would discard trained weights." + ) + # No checkpoint — register fresh adapter logger.debug(f"Rank {self.rank}: Auto-registering new adapter '{model_id}' (no previous checkpoint found)") - self._register_fresh_adapter(model_id) + default_lr = self._ensure_fresh_adapter_materialized(model_id) + self.broadcast_adapter_state(model_id, default_lr) return True, None # Auto-load from checkpoint 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 - self.trainer.register_lora_adapter(model_id, effective_lr) - - if self.rank == 0: - self.trainer.load_adapter_state( - model_id=model_id, - path=checkpoint_path, - load_optimizer=True, - ) - - self.broadcast_adapter_state(model_id, effective_lr) + had_session_spec = model_id in self.trainer.lora_session_specs + if not had_session_spec: + local_error = None + try: + session_spec = load_session_spec_from_checkpoint(checkpoint_path) + self.trainer.register_session(model_id=model_id, session_spec=session_spec, materialize=False) + except Exception as e: + local_error = f"Failed to register session spec from checkpoint for model_id={model_id}: {e}" + + synced_error = self._sync_collective_error(local_error) + if synced_error: + raise RuntimeError(synced_error) + + session_spec = self.trainer.get_lora_session_spec(model_id) + effective_lr = session_spec["optimizer_config"]["learning_rate"] + created_adapter = self._ensure_adapter_materialized_for_restore(model_id, effective_lr) + self._restore_adapter_state( + model_id=model_id, + path=checkpoint_path, + load_optimizer=True, + lr=None, + default_lr=effective_lr, + created_adapter=created_adapter, + ) adapter_state = self.trainer.adapter_manager.get_adapter_state(model_id) logger.debug( @@ -191,12 +684,59 @@ def auto_load_if_evicted(self, model_id: str) -> Tuple[bool, Optional[str]]: f"Rank {self.rank}: Failed to auto-load adapter '{model_id}' from {checkpoint_path}: {e}", exc_info=True, ) - return False, None + raise RuntimeError(f"Failed to auto-load adapter '{model_id}' from {checkpoint_path}: {e}") from e # ======================================================================== # Adapter Registration Handler # ======================================================================== + async def handle_register_session(self, command_dict: Dict[str, Any]) -> Dict[str, Any]: + """Handle register session request (all ranks participate).""" + p: RegisterSessionData = command_dict.get("payload", RegisterSessionData()) + model_id = p.model_id + session_spec = p.session_spec + materialize = p.materialize + had_session_spec = model_id in self.trainer.lora_session_specs + had_adapter = self.trainer.adapter_manager is not None and self.trainer.adapter_manager.has_adapter(model_id) + local_error = None + result = None + + logger.debug( + f"Rank {self.rank}: Registering session: model_id={model_id}, " + f"materialize={materialize}, session_spec={session_spec}" + ) + + try: + result = self.trainer.register_session( + model_id=model_id, + session_spec=session_spec, + materialize=materialize, + ) + except Exception as e: + logger.error(f"Rank {self.rank}: register_session failed: {e}", exc_info=True) + local_error = str(e) + + synced_error = self._sync_collective_error(local_error) + if synced_error: + self._rollback_session_registration( + model_id, + had_session_spec=had_session_spec, + had_adapter=had_adapter, + ) + raise RuntimeError(f"Session registration failed: {synced_error}") + + if ( + materialize + and self.trainer.adapter_manager is not None + and self.trainer.adapter_manager.has_adapter(model_id) + ): + default_lr = session_spec["optimizer_config"]["learning_rate"] + self.broadcast_adapter_state(model_id, default_lr) + + if self.rank == 0: + return result + return {} + async def handle_register_adapter(self, command_dict: Dict[str, Any]) -> Dict[str, Any]: """ Handle register adapter request (all ranks participate). @@ -207,28 +747,38 @@ async def handle_register_adapter(self, command_dict: Dict[str, Any]) -> Dict[st Returns: Dict with registration result """ - try: - p: RegisterAdapterData = command_dict.get("payload", RegisterAdapterData()) - model_id = p.model_id - lr = p.lr + p: RegisterAdapterData = command_dict.get("payload", RegisterAdapterData()) + model_id = p.model_id + lr = p.lr + had_session_spec = model_id in self.trainer.lora_session_specs + had_adapter = self.trainer.adapter_manager is not None and self.trainer.adapter_manager.has_adapter(model_id) + local_error = None + result = None - logger.debug(f"Rank {self.rank}: Registering adapter: model_id={model_id}, lr={lr}") + logger.debug(f"Rank {self.rank}: Registering adapter: model_id={model_id}, lr={lr}") + try: result = self.trainer.register_adapter(model_id=model_id, lr=lr) + except Exception as e: + logger.error(f"Rank {self.rank}: register_adapter failed: {e}", exc_info=True) + local_error = str(e) + + synced_error = self._sync_collective_error(local_error) + if synced_error: + self._rollback_session_registration( + model_id, + had_session_spec=had_session_spec, + had_adapter=had_adapter, + ) + raise RuntimeError(f"Adapter registration failed: {synced_error}") - logger.debug(f"Rank {self.rank}: register_adapter completed: model_id={model_id}") + self.broadcast_adapter_state(model_id, lr) - if self.rank == 0: - return result - else: - return {} + logger.debug(f"Rank {self.rank}: register_adapter completed: model_id={model_id}") - except Exception as e: - logger.error(f"Rank {self.rank}: register_adapter failed: {e}", exc_info=True) - return { - "registered": False, - "error": f"Adapter registration failed: {str(e)}", - } + if self.rank == 0: + return result + return {} # ======================================================================== # Adapter State Save/Load Handlers @@ -255,6 +805,9 @@ async def handle_save_adapter_state(self, command_dict: Dict[str, Any]) -> Dict[ logger.debug(f"Rank {self.rank}: Participating in save_adapter_state: model_id={model_id}, path={path}") + if self.trainer.adapter_manager is not None: + self.auto_load_if_evicted(model_id, allow_fresh_materialization=False) + result = self.trainer.save_adapter_state( model_id=model_id, path=path, @@ -270,17 +823,15 @@ async def handle_save_adapter_state(self, command_dict: Dict[str, Any]) -> Dict[ except Exception as e: logger.error(f"Rank {self.rank}: save_adapter_state failed: {e}", exc_info=True) - return { - "success": False, - "error": f"Adapter state save failed: {str(e)}", - } + raise RuntimeError(f"Adapter state save failed: {str(e)}") from e async def handle_load_adapter_state(self, command_dict: Dict[str, Any]) -> Dict[str, Any]: """ Handle load adapter state request. - ALL ranks must register the adapter, but only rank 0 loads weights - from disk, then broadcasts to other ranks. + Restores adapter state according to `lora.adapter_state_load_mode`: + either all ranks read the checkpoint locally, or rank 0 reads it and + broadcasts weights, metadata, and optimizer state. Args: command_dict: Command dictionary with model_id, path, load_optimizer, lr @@ -288,40 +839,58 @@ async def handle_load_adapter_state(self, command_dict: Dict[str, Any]) -> Dict[ Returns: Dict with load result (from rank 0 only) """ + had_session_spec = False + had_adapter = False + model_id = "" try: p: AdapterStateData = command_dict.get("payload", AdapterStateData()) model_id = p.model_id path = p.path load_optimizer = p.load_optimizer lr = p.lr + had_session_spec = model_id in self.trainer.lora_session_specs + had_adapter = self.trainer.adapter_manager is not None and self.trainer.adapter_manager.has_adapter( + model_id + ) if not path: raise ValueError("path is required for load_adapter_state") - # Step 1: ALL ranks register the adapter with fresh weights - effective_lr = lr if lr is not None else 1e-5 - if not self.trainer.adapter_manager.has_adapter(model_id): - logger.debug(f"Rank {self.rank}: Registering adapter for model_id={model_id}") - self.trainer.register_lora_adapter(model_id, effective_lr) + local_error = None + if not had_session_spec: + try: + session_spec = load_session_spec_from_checkpoint(path) + self.trainer.register_session(model_id=model_id, session_spec=session_spec, materialize=False) + except Exception as e: + local_error = f"Failed to register session spec from checkpoint for model_id={model_id}: {e}" - # Step 2: Only rank 0 loads weights and optimizer from disk - result = None - if self.rank == 0: - logger.debug(f"Rank {self.rank}: Loading adapter weights from disk: model_id={model_id}, path={path}") + synced_error = self._sync_collective_error(local_error) + if synced_error: + raise RuntimeError(synced_error) - result = self.trainer.load_adapter_state( - model_id=model_id, - path=path, - load_optimizer=load_optimizer, - lr=lr, - ) + session_spec = self.trainer.get_lora_session_spec(model_id) + # Step 1: ALL ranks register the adapter with fresh weights + effective_lr = lr if lr is not None else session_spec["optimizer_config"]["learning_rate"] + created_adapter = self._ensure_adapter_materialized_for_restore(model_id, effective_lr) - logger.debug( - f"Rank {self.rank}: load_adapter_state completed: model_id={model_id}, step={result.get('step', 0)}" - ) + # Step 2: Restore weights + metadata, and optimizer state if requested, + # using the configured adapter_state_load_mode. + logger.debug( + f"Rank {self.rank}: Restoring adapter state from disk: " + f"model_id={model_id}, path={path}, mode={self._get_adapter_state_load_mode()}" + ) + result = self._restore_adapter_state( + model_id=model_id, + path=path, + load_optimizer=load_optimizer, + lr=lr, + default_lr=effective_lr, + created_adapter=created_adapter, + ) - # Step 3: Broadcast loaded weights from rank 0 to all other ranks - self.broadcast_adapter_state(model_id, effective_lr) + logger.debug( + f"Rank {self.rank}: load_adapter_state completed: model_id={model_id}, step={result.get('step', 0)}" + ) if self.rank == 0: return result @@ -329,6 +898,11 @@ async def handle_load_adapter_state(self, command_dict: Dict[str, Any]) -> Dict[ return {"success": True, "model_id": model_id} except Exception as e: + self._rollback_session_registration( + model_id, + had_session_spec=had_session_spec, + had_adapter=had_adapter, + ) logger.error(f"Rank {self.rank}: load_adapter_state failed: {e}", exc_info=True) return { "success": False, @@ -385,9 +959,9 @@ async def handle_kill_session(self, command_dict: Dict[str, Any]) -> Dict[str, A """ Handle kill session request (full-weights training only). - In full-weights training mode (enable_lora=False), the server operates in - single-tenant mode. This kills the active session to allow starting a new one. - For LoRA mode, this is a no-op since multi-tenancy is supported. + In LoRA mode this removes the resident adapter and the registered + session spec from workers. In full-weights mode it resets the + single active session. Args: command_dict: Command dictionary with model_id and save_checkpoint @@ -403,19 +977,20 @@ async def handle_kill_session(self, command_dict: Dict[str, Any]) -> Dict[str, A f"Rank {self.rank}: Handling kill_session for model_id={model_id}, save_checkpoint={save_checkpoint}" ) + local_error = None + result = None try: result = self.trainer.kill_session(model_id=model_id, save_checkpoint=save_checkpoint) + except Exception as e: + logger.error(f"Rank {self.rank}: kill_session failed: {e}", exc_info=True) + local_error = str(e) - if self.world_size > 1: - dist.barrier() + synced_error = self._sync_collective_error(local_error) + if synced_error: + raise RuntimeError(f"Kill session failed: {synced_error}") - logger.debug(f"Rank {self.rank}: kill_session completed: {result}") - return result + if self.world_size > 1: + dist.barrier() - except Exception as e: - logger.error(f"Rank {self.rank}: kill_session failed: {e}", exc_info=True) - return { - "success": False, - "message": f"Kill session failed: {str(e)}", - "checkpoint_path": None, - } + logger.debug(f"Rank {self.rank}: kill_session completed: {result}") + return result diff --git a/src/xorl/server/runner/adapters/manager.py b/src/xorl/server/runner/adapters/manager.py index 40627367..d4ca5afc 100644 --- a/src/xorl/server/runner/adapters/manager.py +++ b/src/xorl/server/runner/adapters/manager.py @@ -17,15 +17,25 @@ import math import os import time +from copy import deepcopy from dataclasses import dataclass, field -from typing import Any, Dict, List, Optional +from typing import Any, Dict, List, Optional, Tuple import torch import torch.nn as nn 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 +from xorl.lora.utils import ( + convert_peft_lora_state_dict, + get_lora_tensor_shard_specs, +) +from xorl.optim import build_optimizer +from xorl.server.session_spec import ( + load_session_spec_from_checkpoint, + session_optimizer_build_kwargs, + write_session_spec, +) try: @@ -50,6 +60,7 @@ class AdapterState: """ model_id: str + session_spec: Dict[str, Any] lora_params: Dict[str, nn.Parameter] # Actual Parameters with own .grad optimizer: torch.optim.Optimizer # Per-adapter optimizer global_step: int = 0 @@ -80,6 +91,14 @@ def __init__( checkpoint_dir: Optional[str] = None, auto_save_on_eviction: bool = True, lora_config: Optional[Dict[str, Any]] = None, + optimizer_config: Optional[Dict[str, Any]] = None, + optimizer_type: str = "adamw", + optimizer_dtype: str = "bf16", + optimizer_kwargs: Optional[Dict[str, Any]] = None, + weight_decay: float = 0.01, + betas: Tuple[float, float] = (0.9, 0.95), + eps: float = 1e-8, + optimizer_fused: Optional[bool] = None, ): """ Initialize the adapter manager. @@ -91,6 +110,14 @@ def __init__( checkpoint_dir: Directory for saving adapter checkpoints (default: outputs/adapters) auto_save_on_eviction: If True, save adapter state before LRU eviction lora_config: LoRA configuration dict (for saving adapter_config.json) + optimizer_config: Training optimizer configuration for per-adapter optimizers + optimizer_type: Optimizer type passed to xorl.optim.build_optimizer + optimizer_dtype: Optimizer state dtype for supported optimizers + optimizer_kwargs: Optimizer-specific kwargs (e.g. Muon settings) + weight_decay: Weight decay used when building adapter optimizers + betas: Beta coefficients for Adam-family optimizers + eps: Epsilon used by Adam-family optimizers + optimizer_fused: Whether to request fused optimizer kernels """ self.model = model self.device = device @@ -98,24 +125,327 @@ def __init__( self.checkpoint_dir = checkpoint_dir or "outputs/adapters" self.auto_save_on_eviction = auto_save_on_eviction self.lora_config = lora_config or {} + self.optimizer_config = optimizer_config or {} + self.optimizer_type = optimizer_type + self.optimizer_dtype = optimizer_dtype + self.optimizer_kwargs = deepcopy(optimizer_kwargs or {}) + self.weight_decay = weight_decay + self.betas = betas + self.eps = eps + self.optimizer_fused = device.type == "cuda" if optimizer_fused is None else optimizer_fused self.adapters: Dict[str, AdapterState] = {} self.current_adapter_id: Optional[str] = None # Cache the list of LoRA parameter names for efficient lookups self._lora_param_names: List[str] = [] + self._lora_param_metadata: Dict[str, Dict[str, Any]] = {} for name, param in self.model.named_parameters(): if "lora_A" in name or "lora_B" in name: self._lora_param_names.append(name) + param_shape = tuple(param.shape if _HAS_DTENSOR and isinstance(param, DTensor) else param.data.shape) + self._lora_param_metadata[name] = { + "shape": param_shape, + "dtype": param.dtype if _HAS_DTENSOR and isinstance(param, DTensor) else param.data.dtype, + "rank_dim": self._infer_lora_rank_dim(name, param_shape), + } logger.info( f"LoRAAdapterManager initialized with {len(self._lora_param_names)} LoRA parameters, " - f"max_adapters={max_adapters}, auto_save_on_eviction={auto_save_on_eviction}" + f"max_adapters={max_adapters}, auto_save_on_eviction={auto_save_on_eviction}, " + f"optimizer={optimizer_type}" ) + @staticmethod + def _infer_lora_rank_dim(name: str, shape: Tuple[int, ...]) -> int: + """Infer which tensor dimension corresponds to the LoRA rank.""" + if "lora_A" in name: + if len(shape) == 2: + return 0 + if len(shape) == 3: + return 2 + if "lora_B" in name: + if len(shape) == 2: + return 1 + if len(shape) == 3: + return 1 + raise ValueError(f"Cannot infer LoRA rank dimension for parameter {name!r} with shape {shape!r}") + + @staticmethod + def _replace_dim(shape: Tuple[int, ...], dim: int, value: int) -> Tuple[int, ...]: + updated = list(shape) + updated[dim] = value + return tuple(updated) + + @staticmethod + def _slice_to_rank(tensor: torch.Tensor, *, rank_dim: int, active_rank: int) -> torch.Tensor: + return tensor.narrow(rank_dim, 0, active_rank) + + @staticmethod + def _slice_to_shape(tensor: torch.Tensor, *, rank_dim: int, target_shape: Tuple[int, ...]) -> torch.Tensor: + active_rank = target_shape[rank_dim] + sliced = tensor.narrow(rank_dim, 0, active_rank) + if tuple(sliced.shape) != target_shape: + raise ValueError(f"Expected sliced tensor shape {target_shape}, got {tuple(sliced.shape)}") + return sliced + + @staticmethod + def _expand_compact_tensor( + tensor: torch.Tensor, + *, + full_shape: Tuple[int, ...], + rank_dim: int, + ) -> torch.Tensor: + if tuple(tensor.shape) == full_shape: + return tensor + expanded = torch.zeros(full_shape, dtype=tensor.dtype, device=tensor.device) + slices = [slice(None)] * len(full_shape) + slices[rank_dim] = slice(0, tensor.shape[rank_dim]) + expanded[tuple(slices)] = tensor + return expanded + + @staticmethod + def _session_rank(session_spec: Dict[str, Any]) -> int: + return int(session_spec["lora_config"]["lora_rank"]) + + @staticmethod + def _session_alpha(session_spec: Dict[str, Any]) -> int: + return int(session_spec["lora_config"]["lora_alpha"]) + + @staticmethod + def _strip_optimizer_config(session_spec: Dict[str, Any]) -> Dict[str, Any]: + """Return the structural part of a LoRA session spec without optimizer metadata.""" + stripped = deepcopy(session_spec) + stripped.pop("optimizer_config", None) + return stripped + + @staticmethod + def _strip_optimizer_learning_rate(session_spec: Dict[str, Any]) -> Dict[str, Any]: + """Return a session spec without the mutable optimizer learning-rate field.""" + stripped = deepcopy(session_spec) + optimizer_config = stripped.get("optimizer_config") + if isinstance(optimizer_config, dict): + optimizer_config.pop("learning_rate", None) + return stripped + + @staticmethod + def _serialize_optimizer_metadata_value(value: Any) -> Any: + """Convert optimizer metadata into JSON-safe values.""" + if isinstance(value, torch.dtype): + if value == torch.bfloat16: + return "bf16" + if value == torch.float32: + return "fp32" + return str(value) + if isinstance(value, dict): + return {k: LoRAAdapterManager._serialize_optimizer_metadata_value(v) for k, v in value.items()} + if isinstance(value, (list, tuple)): + return [LoRAAdapterManager._serialize_optimizer_metadata_value(v) for v in value] + return value + + @staticmethod + def _update_state_learning_rate(state: AdapterState, lr: float) -> None: + """Keep adapter LR, optimizer param groups, and session spec in sync.""" + state.lr = float(lr) + state.session_spec.setdefault("optimizer_config", {})["learning_rate"] = state.lr + for param_group in state.optimizer.param_groups: + if state.session_spec.get("optimizer_config", {}).get("type") == "muon" and param_group.get( + "use_muon", False + ): + continue + param_group["lr"] = state.lr + + def _max_supported_session_rank(self) -> int: + """Return the largest LoRA rank the live model substrate can support.""" + if not self._lora_param_metadata: + raise RuntimeError("Cannot determine LoRA rank capacity: model does not expose any LoRA parameters.") + return min(metadata["shape"][metadata["rank_dim"]] for metadata in self._lora_param_metadata.values()) + + def _validate_session_rank_against_model_capacity(self, session_spec: Dict[str, Any]) -> None: + """Reject session specs whose runtime rank exceeds the live model capacity.""" + session_rank = self._session_rank(session_spec) + max_supported_rank = self._max_supported_session_rank() + if session_rank > max_supported_rank: + raise ValueError( + f"Session rank {session_rank} exceeds live model LoRA capacity {max_supported_rank}. " + "Restart the server with a larger max_lora_rank-compatible model substrate before loading this checkpoint." + ) + + @staticmethod + def _module_name_for_lora_param(name: str) -> str: + """Extract the target module name from an internal LoRA parameter name.""" + base_name = ( + name.replace(".lora_A.weight", "") + .replace(".lora_B.weight", "") + .replace(".lora_A", "") + .replace(".lora_B", "") + .replace("_lora_A", "") + .replace("_lora_B", "") + ) + parts = base_name.split(".") + if not parts: + raise ValueError(f"Cannot infer target module from LoRA parameter name {name!r}") + return parts[-1] + + @staticmethod + def _canonical_lora_param_name(name: str) -> str: + """Normalize LoRA parameter names across checkpoint formats.""" + if name.endswith(".weight"): + return name[: -len(".weight")] + return name + + def _expected_target_modules(self) -> List[str]: + """Return the live model's expected LoRA target modules.""" + return sorted( + { + self._module_name_for_lora_param(name) + for name in self._lora_param_names + if "lora_A" in name or "lora_B" in name + } + ) + + def _validate_checkpoint_adapter_config(self, path: str) -> None: + """Validate checkpoint-level adapter structure against the live model configuration.""" + adapter_config_path = os.path.join(path, "adapter_config.json") + if not os.path.exists(adapter_config_path): + return + + with open(adapter_config_path, "r") as f: + adapter_config = json.load(f) + + checkpoint_target_modules = adapter_config.get("target_modules") + if checkpoint_target_modules is not None: + actual_target_modules = sorted(str(module) for module in checkpoint_target_modules) + expected_target_modules = self._expected_target_modules() + if actual_target_modules != expected_target_modules: + raise ValueError( + "Checkpoint target_modules do not match the live LoRA adapter structure. " + f"checkpoint={actual_target_modules!r}, live={expected_target_modules!r}" + ) + + if "moe_hybrid_shared_lora" in adapter_config: + checkpoint_hybrid = bool(adapter_config["moe_hybrid_shared_lora"]) + expected_hybrid = bool(self.lora_config.get("moe_hybrid_shared_lora", False)) + if checkpoint_hybrid != expected_hybrid: + raise ValueError( + "Checkpoint moe_hybrid_shared_lora does not match the live LoRA adapter structure. " + f"checkpoint={checkpoint_hybrid!r}, live={expected_hybrid!r}" + ) + + def get_optimizer_metadata(self) -> Dict[str, Any]: + """Return a JSON-safe description of the adapter optimizer contract.""" + return { + "type": self.optimizer_type, + "dtype": self.optimizer_dtype, + "weight_decay": self.weight_decay, + "betas": list(self.betas), + "eps": self.eps, + "optimizer_kwargs": self._serialize_optimizer_metadata_value(self.optimizer_kwargs), + } + + def get_adapter_session_spec(self, model_id: str) -> Dict[str, Any]: + """Return the normalized session spec for an adapter.""" + return deepcopy(self.get_adapter_state(model_id).session_spec) + + def _legacy_session_spec(self, *, lr: float) -> Dict[str, Any]: + """Build a session spec for compatibility call sites that only provide lr.""" + default_rank = self.lora_config.get("lora_rank") + if default_rank is None and self._lora_param_names: + metadata = self._lora_param_metadata[self._lora_param_names[0]] + default_rank = metadata["shape"][metadata["rank_dim"]] + default_alpha = self.lora_config.get("lora_alpha", default_rank or 16) + # Start from the manager-level optimizer_config so passthrough flags + # like cautious_weight_decay reach build_optimizer; structured fields + # below override anything the manager-level dict supplies. + optimizer_config: Dict[str, Any] = dict(self.optimizer_config or {}) + weight_decay = optimizer_config.get("weight_decay", self.weight_decay) + optimizer_config.update( + { + "type": self.optimizer_type, + "learning_rate": float(lr), + "weight_decay": float(weight_decay), + "optimizer_dtype": self.optimizer_dtype, + "betas": list(self.betas), + "eps": float(self.eps), + "optimizer_kwargs": self._serialize_optimizer_metadata_value(self.optimizer_kwargs), + } + ) + return { + "base_model": self.lora_config.get("base_model", ""), + "is_lora": True, + "lora_config": { + "lora_rank": int(default_rank or 32), + "lora_alpha": int(default_alpha), + }, + "optimizer_config": optimizer_config, + } + + def _set_model_runtime_lora_config(self, *, lora_rank: int, lora_alpha: int) -> None: + """Update all model-side LoRA modules to use the active session rank/alpha.""" + for module in self.model.modules(): + setter = getattr(module, "set_runtime_lora_config", None) + if setter is not None: + setter(lora_rank, lora_alpha) + + @staticmethod + def _build_parameter_module(lora_params: Dict[str, nn.Parameter]) -> nn.Module: + """Wrap an adapter's parameters in a temporary module with stable parameter names.""" + root = nn.Module() + for full_name, param in lora_params.items(): + current = root + parts = full_name.split(".") + for part in parts[:-1]: + child = current._modules.get(part) + if child is None: + child = nn.Module() + current.add_module(part, child) + current = child + + leaf_name = parts[-1] + if leaf_name in current._parameters: + raise ValueError(f"Duplicate parameter name while building adapter optimizer module: {full_name}") + current.register_parameter(leaf_name, param) + return root + + def _build_adapter_optimizer(self, lora_params: Dict[str, nn.Parameter], lr: float) -> torch.optim.Optimizer: + """Build an optimizer for one adapter via the shared optimizer factory.""" + adapter_module = self._build_parameter_module(lora_params) + return build_optimizer( + adapter_module, + lr=lr, + betas=self.betas, + eps=self.eps, + weight_decay=self.weight_decay, + fused=self.optimizer_fused, + optimizer_type=self.optimizer_type, + optimizer_dtype=self.optimizer_dtype, + optimizer_kwargs=deepcopy(self.optimizer_kwargs), + ) + + def _build_adapter_optimizer_for_session( + self, lora_params: Dict[str, nn.Parameter], session_spec: Dict[str, Any] + ) -> torch.optim.Optimizer: + adapter_module = self._build_parameter_module(lora_params) + build_kwargs = session_optimizer_build_kwargs(session_spec["optimizer_config"]) + return build_optimizer( + adapter_module, + fused=self.optimizer_fused, + **build_kwargs, + ) + + @staticmethod + def _has_pending_gradients(state: AdapterState) -> bool: + """Return whether an adapter has captured gradients awaiting an optimizer step.""" + return any(param.grad is not None for param in state.lora_params.values()) + def _maybe_evict(self) -> Optional[str]: """ Evict the least recently used adapter if at capacity. + Adapters with pending gradients are not evictable because checkpointing + them would silently drop the captured gradients before `optim_step`. + If every resident adapter has pending gradients, this raises instead of + discarding training state. + If auto_save_on_eviction is enabled, saves the adapter state before evicting. Returns: @@ -124,8 +454,17 @@ def _maybe_evict(self) -> Optional[str]: if len(self.adapters) >= self.max_adapters: if not self.adapters: return None - # Find LRU adapter - all adapters can be evicted - lru_id = min(self.adapters.keys(), key=lambda k: self.adapters[k].last_access_time) + evictable_ids = [ + model_id for model_id, state in self.adapters.items() if not self._has_pending_gradients(state) + ] + if not evictable_ids: + raise RuntimeError( + "Cannot evict any adapter safely because all resident adapters have pending gradients. " + "Call optim_step for at least one session before loading or creating another adapter." + ) + + # Find the LRU adapter among the clean (step-complete) adapters. + lru_id = min(evictable_ids, key=lambda k: self.adapters[k].last_access_time) logger.info(f"Evicting LRU adapter: {lru_id} (capacity {len(self.adapters)}/{self.max_adapters})") # Auto-save before eviction if enabled @@ -144,7 +483,8 @@ def _maybe_evict(self) -> Optional[str]: def register_adapter( self, model_id: str, - lr: float, + lr: Optional[float] = None, + session_spec: Optional[Dict[str, Any]] = None, initialize_fresh: bool = True, ) -> None: """ @@ -155,10 +495,27 @@ def register_adapter( Args: model_id: Unique identifier for this training run - lr: Learning rate for this adapter's optimizer + lr: Optional learning rate override for legacy call sites + session_spec: Normalized session runtime spec for this adapter initialize_fresh: If True, initialize with fresh random weights. If False, use the current model's LoRA weights. """ + effective_lr = float(lr) if lr is not None else None + if session_spec is None: + if effective_lr is None: + effective_lr = 1e-5 + session_spec = self._legacy_session_spec(lr=effective_lr) + else: + session_spec = deepcopy(session_spec) + if effective_lr is not None: + session_spec["optimizer_config"]["learning_rate"] = effective_lr + + self._validate_session_rank_against_model_capacity(session_spec) + session_rank = self._session_rank(session_spec) + session_alpha = self._session_alpha(session_spec) + optimizer_config = session_spec["optimizer_config"] + effective_lr = float(optimizer_config["learning_rate"]) + # Evict LRU adapter if at capacity and this is a new adapter if model_id not in self.adapters: self._maybe_evict() @@ -171,31 +528,24 @@ def register_adapter( lora_params: Dict[str, nn.Parameter] = {} for name, param in self.model.named_parameters(): if name in self._lora_param_names: - # Get shape and dtype from the parameter - # IMPORTANT: For DTensors, use .shape (global shape) and .dtype directly - # DO NOT call full_tensor() here as it's a collective operation that - # requires all ranks to participate, which can cause deadlock when - # called from load_adapter_state (only rank 0 does the load) - if _HAS_DTENSOR and isinstance(param, DTensor): - # DTensor.shape gives the global (unsharded) shape - param_shape = param.shape - param_dtype = param.dtype - else: - param_shape = param.data.shape - param_dtype = param.data.dtype + metadata = self._lora_param_metadata[name] + param_shape = metadata["shape"] + param_dtype = metadata["dtype"] + rank_dim = metadata["rank_dim"] + compact_shape = self._replace_dim(param_shape, rank_dim, session_rank) if initialize_fresh: - # Fresh initialization - create regular tensor on device + # Fresh initialization - create compact regular tensor on device if "lora_A" in name: new_tensor = torch.empty( - param_shape, + compact_shape, dtype=param_dtype, device=self.device, ) nn.init.kaiming_uniform_(new_tensor, a=math.sqrt(5)) else: # lora_B new_tensor = torch.zeros( - param_shape, + compact_shape, dtype=param_dtype, device=self.device, ) @@ -208,32 +558,37 @@ def register_adapter( param_data = param.full_tensor() else: param_data = param.data - new_tensor = param_data.detach().clone().to(self.device) + new_tensor = ( + self._slice_to_shape( + param_data.detach(), + rank_dim=rank_dim, + target_shape=compact_shape, + ) + .clone() + .to(self.device) + ) # Create as nn.Parameter so it has its own .grad slot lora_params[name] = nn.Parameter(new_tensor, requires_grad=True) - # Create optimizer for this adapter's params - optimizer = torch.optim.AdamW( - list(lora_params.values()), - lr=lr, - betas=(0.9, 0.95), - eps=1e-8, - weight_decay=0.01, - ) + # Build optimizer for this adapter using the session's optimizer contract. + optimizer = self._build_adapter_optimizer_for_session(lora_params, session_spec) self.adapters[model_id] = AdapterState( model_id=model_id, + session_spec=session_spec, lora_params=lora_params, optimizer=optimizer, global_step=0, global_forward_backward_step=0, - lr=lr, + lr=effective_lr, ) logger.info( f"Registered adapter for model_id={model_id} " - f"(lr={lr}, fresh_weights={initialize_fresh}, num_params={len(lora_params)})" + f"(rank={session_rank}, alpha={session_alpha}, lr={effective_lr}, " + f"fresh_weights={initialize_fresh}, num_params={len(lora_params)}, " + f"optimizer={optimizer_config['type']})" ) def prepare_forward(self, model_id: str) -> None: @@ -262,13 +617,22 @@ def prepare_forward(self, model_id: str) -> None: state = self.adapters[model_id] # Update last access time for LRU tracking state.last_access_time = time.time() + self._set_model_runtime_lora_config( + lora_rank=self._session_rank(state.session_spec), + lora_alpha=self._session_alpha(state.session_spec), + ) # Copy adapter weights into model's params (for forward to use) # Use no_grad to avoid autograd issues with DTensor views with torch.no_grad(): for name, param in self.model.named_parameters(): if name in state.lora_params: - adapter_data = state.lora_params[name].data + metadata = self._lora_param_metadata[name] + adapter_data = self._expand_compact_tensor( + state.lora_params[name].data, + full_shape=metadata["shape"], + rank_dim=metadata["rank_dim"], + ) if _HAS_DTENSOR and isinstance(param, DTensor): # For DTensor: copy to the local tensor (the shard) @@ -328,10 +692,16 @@ def capture_gradients(self, model_id: str) -> None: if name in state.lora_params: adapter_param = state.lora_params[name] if param.grad is not None: + metadata = self._lora_param_metadata[name] # Handle DTensor (FSDP2 sharded gradients) grad = param.grad if _HAS_DTENSOR and isinstance(grad, DTensor): grad = grad.full_tensor() + grad = self._slice_to_shape( + grad, + rank_dim=metadata["rank_dim"], + target_shape=tuple(adapter_param.shape), + ) # Copy gradient to adapter's param (accumulate for grad accumulation) if adapter_param.grad is None: @@ -372,9 +742,7 @@ def optim_step( state = self.adapters[model_id] # Update learning rate - state.lr = lr - for pg in state.optimizer.param_groups: - pg["lr"] = lr + self._update_state_learning_rate(state, lr) # Deferred gradient normalization: scale raw gradients by 1/accumulated_valid_tokens if accumulated_valid_tokens > 0: @@ -447,9 +815,7 @@ def get_lr(self, model_id: str) -> float: def set_lr(self, model_id: str, lr: float) -> None: """Set the learning rate for an adapter.""" state = self.adapters[model_id] - state.lr = lr - for param_group in state.optimizer.param_groups: - param_group["lr"] = lr + self._update_state_learning_rate(state, lr) def has_adapter(self, model_id: str) -> bool: """Check if an adapter is registered for a model_id.""" @@ -558,10 +924,17 @@ def save_adapter_state( # 1. Save LoRA weights in safetensors format (PEFT-compatible) # Convert parameter names to PEFT format: base_model.model.{name} + raw_weights = {name: param.data.detach() for name, param in state.lora_params.items()} + # Adapter-owned tensors are already compacted to that session's rank. + # Do not slice against the live model: LRU eviction can save a different + # adapter than the one currently loaded into the model scratch space. + active_weights = raw_weights weights_dict = {} - for name, param in state.lora_params.items(): - peft_name = f"base_model.model.{name}" - weights_dict[peft_name] = param.data.cpu().to(torch.bfloat16) + for name, tensor in active_weights.items(): + peft_name = f"base_model.model.{self._canonical_lora_param_name(name)}" + if peft_name in weights_dict: + raise ValueError(f"Duplicate canonical LoRA parameter name while saving adapter state: {peft_name}") + weights_dict[peft_name] = tensor.detach().cpu().contiguous() weights_path = os.path.join(path, "adapter_model.safetensors") safetensors_save_file(weights_dict, weights_path) @@ -571,6 +944,11 @@ def save_adapter_state( optimizer_path = os.path.join(path, "optimizer.pt") torch.save(state.optimizer.state_dict(), optimizer_path) + # 3. Save normalized session runtime spec with the current learning rate. + checkpoint_session_spec = deepcopy(state.session_spec) + checkpoint_session_spec["optimizer_config"]["learning_rate"] = float(state.lr) + write_session_spec(path, checkpoint_session_spec) + # 3. Save metadata metadata = { "model_id": model_id, @@ -579,26 +957,28 @@ def save_adapter_state( "lr": state.lr, "timestamp": time.time(), "save_optimizer": save_optimizer, + "optimizer": deepcopy(checkpoint_session_spec["optimizer_config"]), } metadata_path = os.path.join(path, "metadata.json") with open(metadata_path, "w") as f: json.dump(metadata, f, indent=2) # 4. Save adapter config (PEFT-compatible) - # Extract LoRA config from parameter shapes - lora_r = None target_modules = set() - for name, param in state.lora_params.items(): - if "lora_A" in name: - lora_r = param.shape[0] # lora_A is [r, in_features] + for name, tensor in active_weights.items(): + if "lora_A" in name or "_lora_A" in name: + if name.endswith("_lora_A"): + target_modules.add(name.rsplit(".", 1)[-1][: -len("_lora_A")]) + continue # Extract module name (e.g., "model.layers.0.self_attn.q_proj" from full name) - parts = name.replace(".lora_A.weight", "").split(".") + parts = name.replace(".lora_A.weight", "").replace(".lora_A", "").replace("_lora_A", "").split(".") if len(parts) >= 1: target_modules.add(parts[-1]) # e.g., "q_proj" adapter_config = { - "r": lora_r, - "lora_alpha": lora_r, # Assume alpha = r (common default) + "base_model_name_or_path": state.session_spec.get("base_model"), + "r": self._session_rank(state.session_spec), + "lora_alpha": self._session_alpha(state.session_spec), "target_modules": list(target_modules), "lora_dropout": 0.0, "bias": "none", @@ -660,52 +1040,137 @@ def load_adapter_state( # Determine learning rate effective_lr = lr if lr is not None else metadata.get("lr", 1e-5) + registered_state = self.adapters.get(model_id) + expected_session_spec = deepcopy(registered_state.session_spec) if registered_state is not None else None + if expected_session_spec is None: + expected_session_spec = self._legacy_session_spec(lr=effective_lr) + + checkpoint_session_spec = load_session_spec_from_checkpoint( + path, + fallback_base_model=expected_session_spec.get("base_model"), + fallback_session_spec=expected_session_spec, + ) + self._validate_checkpoint_adapter_config(path) + + if registered_state is not None: + checkpoint_spec_for_compare = checkpoint_session_spec + registered_spec_for_compare = registered_state.session_spec + if lr is not None: + checkpoint_spec_for_compare = self._strip_optimizer_learning_rate(checkpoint_spec_for_compare) + registered_spec_for_compare = self._strip_optimizer_learning_rate(registered_spec_for_compare) + + if load_optimizer: + specs_match = checkpoint_spec_for_compare == registered_spec_for_compare + mismatch_context = "registered multi-adapter session" + else: + specs_match = self._strip_optimizer_config(checkpoint_spec_for_compare) == self._strip_optimizer_config( + registered_spec_for_compare + ) + mismatch_context = "registered multi-adapter session for weights-only restore" - # 2. Register adapter if not exists (this will evict if needed) - if model_id not in self.adapters: - self.register_adapter(model_id, lr=effective_lr, initialize_fresh=True) - - state = self.adapters[model_id] + if not specs_match: + raise ValueError( + "Checkpoint session spec does not match the " + f"{mismatch_context}. checkpoint={checkpoint_session_spec!r}, " + f"current={registered_state.session_spec!r}" + ) - # 3. Load LoRA weights - weights_path = os.path.join(path, "adapter_model.safetensors") - if os.path.exists(weights_path): - loaded_weights = safetensors_load_file(weights_path) - 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}" + # 2. Register adapter if not exists (this will evict if needed). + # Track whether this call did the registration so a downstream load + # failure does not leave a fresh-init adapter resident under model_id. + registered_here = False + if model_id not in self.adapters: + self.register_adapter( + model_id, + session_spec=checkpoint_session_spec, + initialize_fresh=True, + ) + registered_here = True + + try: + state = self.adapters[model_id] + + # 3. Load LoRA weights + weights_path = os.path.join(path, "adapter_model.safetensors") + if os.path.exists(weights_path): + loaded_weights = safetensors_load_file(weights_path) + expected_param_map: Dict[str, str] = {} + expected_shapes: Dict[str, torch.Size] = {} + for actual_name in state.lora_params: + canonical_name = self._canonical_lora_param_name(actual_name) + if canonical_name in expected_param_map and expected_param_map[canonical_name] != actual_name: + raise ValueError( + f"Live adapter contains duplicate LoRA tensors after canonicalization. param={canonical_name!r}" + ) + expected_param_map[canonical_name] = actual_name + expected_shapes[canonical_name] = state.lora_params[actual_name].shape + + expected_shard_specs = get_lora_tensor_shard_specs(self.model, names=expected_shapes.keys()) + converted_weights = convert_peft_lora_state_dict( + loaded_weights, + expected_shapes=expected_shapes, + expected_shard_specs=expected_shard_specs, ) + checkpoint_tensors: Dict[str, torch.Tensor] = {} + for converted_name, weight in converted_weights.items(): + canonical_name = self._canonical_lora_param_name(converted_name) + if canonical_name in checkpoint_tensors: + raise ValueError( + f"Checkpoint contains duplicate LoRA tensors after canonicalization. param={canonical_name!r}" + ) + checkpoint_tensors[canonical_name] = weight.to(self.device) + + expected_param_names = set(expected_param_map) + checkpoint_param_names = set(checkpoint_tensors) + missing_param_names = sorted(expected_param_names - checkpoint_param_names) + unexpected_param_names = sorted(checkpoint_param_names - expected_param_names) + if missing_param_names or unexpected_param_names: + raise ValueError( + "Checkpoint LoRA parameter set does not match the live adapter structure. " + f"missing={missing_param_names!r}, unexpected={unexpected_param_names!r}" + ) + + for internal_name, tensor in checkpoint_tensors.items(): + target_param = state.lora_params[expected_param_map[internal_name]] + if tuple(tensor.shape) != tuple(target_param.shape): + raise ValueError( + "Checkpoint tensor shape does not match the live adapter shape. " + f"param={internal_name!r}, checkpoint={tuple(tensor.shape)!r}, " + f"live={tuple(target_param.shape)!r}" + ) - 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}") + for internal_name, tensor in checkpoint_tensors.items(): + state.lora_params[expected_param_map[internal_name]].data.copy_(tensor) + else: + raise FileNotFoundError(f"Weights file not found: {weights_path}") - # 4. Load optimizer state - optimizer_path = os.path.join(path, "optimizer.pt") - if load_optimizer and os.path.exists(optimizer_path): - optimizer_state = torch.load(optimizer_path, map_location=self.device, weights_only=True) - state.optimizer.load_state_dict(optimizer_state) - logger.debug(f"Loaded optimizer state from {optimizer_path}") + # 4. Load optimizer state + optimizer_path = os.path.join(path, "optimizer.pt") + if load_optimizer and os.path.exists(optimizer_path): + optimizer_state = torch.load(optimizer_path, map_location=self.device, weights_only=True) + state.optimizer.load_state_dict(optimizer_state) + logger.debug(f"Loaded optimizer state from {optimizer_path}") + except Exception: + if registered_here: + try: + self.remove_adapter(model_id) + except Exception as cleanup_error: + logger.warning( + f"Cleanup remove_adapter({model_id}) after failed load_adapter_state raised: {cleanup_error}" + ) + raise # 5. Restore metadata state.global_step = metadata.get("global_step", 0) state.global_forward_backward_step = metadata.get("global_forward_backward_step", 0) if lr is not None: - state.lr = lr - for pg in state.optimizer.param_groups: - pg["lr"] = lr - elif "lr" in metadata: - state.lr = metadata["lr"] + self._update_state_learning_rate(state, lr) + elif "lr" in metadata and ( + load_optimizer + or self._strip_optimizer_learning_rate(checkpoint_session_spec) + == self._strip_optimizer_learning_rate(state.session_spec) + ): + self._update_state_learning_rate(state, metadata["lr"]) # Update last access time state.last_access_time = time.time() diff --git a/src/xorl/server/runner/checkpoint/manager.py b/src/xorl/server/runner/checkpoint/manager.py index f411a31d..7aa170c8 100644 --- a/src/xorl/server/runner/checkpoint/manager.py +++ b/src/xorl/server/runner/checkpoint/manager.py @@ -24,6 +24,7 @@ import torch.nn as nn from safetensors.torch import save_file from torch.distributed._tensor import DTensor +from torch.distributed.checkpoint import FileSystemReader from torch.distributed.checkpoint.state_dict import StateDictOptions, get_model_state_dict from xorl.checkpoint import ckpt_to_state_dict @@ -31,7 +32,9 @@ from xorl.distributed.parallel_state import get_parallel_state from xorl.lora.utils import get_lora_state_dict, save_lora_checkpoint from xorl.models import save_model_weights +from xorl.server.session_spec import write_session_spec from xorl.utils import helper +from xorl.utils.device import get_device_type logger = logging.getLogger(__name__) @@ -107,10 +110,114 @@ def _get_lora_save_config(self): return target_modules, lora_alpha + def _tensor_is_meta(self, tensor: torch.Tensor) -> bool: + raw_tensor = tensor.data if hasattr(tensor, "data") else tensor + if isinstance(raw_tensor, DTensor): + local_tensor = raw_tensor.to_local() + if hasattr(local_tensor, "wait"): + local_tensor = local_tensor.wait() + return getattr(local_tensor, "is_meta", False) + return getattr(raw_tensor, "is_meta", False) + + def _model_has_meta_tensors(self) -> bool: + for parameter in self.model.parameters(): + if self._tensor_is_meta(parameter): + return True + for buffer in self.model.buffers(): + if self._tensor_is_meta(buffer): + return True + return False + + def _materialize_meta_tensors_for_dcp_load(self) -> bool: + if self.train_config.get("load_weights_mode") != "skip": + return False + if not self._model_has_meta_tensors(): + return False + + device_type = get_device_type() + target_device = "cpu" if device_type == "cpu" else f"{device_type}:{self.local_rank}" + logger.info("Materializing meta parameters before DCP load via to_empty(device=%s)", target_device) + self.model.to_empty(device=target_device) + return True + + def _checkpoint_has_optimizer(self, checkpoint_path: str) -> bool: + if not os.path.exists(os.path.join(checkpoint_path, ".metadata")): + return False + try: + metadata = FileSystemReader(checkpoint_path).read_metadata() + except Exception as exc: # noqa: BLE001 + logger.warning("Could not inspect DCP metadata at %s for optimizer state: %s", checkpoint_path, exc) + return False + return any(key.startswith("optimizer") for key in metadata.state_dict_metadata.keys()) + + def _build_dcp_load_state(self, checkpoint_path: str, load_optimizer: bool) -> Dict[str, Any]: + self._materialize_meta_tensors_for_dcp_load() + + state: Dict[str, Any] = {"model": self.model, "extra_state": {}} + if load_optimizer: + if self.optimizer is not None and self._checkpoint_has_optimizer(checkpoint_path): + state["optimizer"] = self.optimizer + else: + logger.info("DCP checkpoint has no optimizer state; loading model weights only.") + return state + # ------------------------------------------------------------------ # Adapter save / load (multi-tenancy LoRA) # ------------------------------------------------------------------ + def _sync_collective_error(self, local_error: Optional[str]) -> Optional[str]: + """Synchronize save failures across ranks before any barrier.""" + if not dist.is_available() or not dist.is_initialized(): + return local_error + + world_size = dist.get_world_size() + if world_size <= 1: + return local_error + + backend = dist.get_backend() + device = ( + torch.device(f"cuda:{self.local_rank}") + if backend == "nccl" and torch.cuda.is_available() + else torch.device("cpu") + ) + has_error = torch.tensor([1 if local_error else 0], dtype=torch.int64, device=device) + dist.all_reduce(has_error, op=dist.ReduceOp.MAX) + + if has_error.item() == 0: + return None + + error_strings = [None] * world_size + dist.all_gather_object(error_strings, local_error or "") + errors = {i: msg for i, msg in enumerate(error_strings) if msg} + if errors: + return "; ".join(f"rank {i}: {msg}" for i, msg in errors.items()) + return local_error + + def _write_adapter_training_artifacts( + self, + path: str, + model_id: str, + adapter_state: Any, + save_optimizer: bool, + ) -> None: + """Write adapter-specific optimizer state and training metadata on rank 0.""" + if save_optimizer: + optimizer_path = os.path.join(path, "optimizer.pt") + torch.save(adapter_state.optimizer.state_dict(), optimizer_path) + + metadata = { + "model_id": model_id, + "global_step": adapter_state.global_step, + "global_forward_backward_step": adapter_state.global_forward_backward_step, + "lr": adapter_state.lr, + "timestamp": time.time(), + "save_optimizer": save_optimizer, + "optimizer": adapter_state.session_spec["optimizer_config"], + } + metadata_path = os.path.join(path, "metadata.json") + with open(metadata_path, "w") as f: + json.dump(metadata, f, indent=2) + def _gather_adapter_lora_params(self, model_id: str) -> Dict[str, torch.Tensor]: """Gather LoRA params from adapter manager with EP support. @@ -132,7 +239,46 @@ def _gather_adapter_lora_params(self, model_id: str) -> Dict[str, torch.Tensor]: return lora_state_dict - def _save_lora_weights(self, save_path: str, model_id: str) -> None: + @staticmethod + def _infer_lora_rank_dim(name: str, tensor: torch.Tensor) -> Optional[int]: + """Infer which tensor dimension stores LoRA rank for dense and MoE LoRA params.""" + if "lora_A" in name: + if tensor.dim() == 3: + return 2 + if tensor.dim() >= 2: + return 0 + if "lora_B" in name: + if tensor.dim() >= 2: + return 1 + return None + + @classmethod + def _slice_lora_state_dict_to_rank( + cls, + lora_state_dict: Dict[str, torch.Tensor], + active_rank: int, + ) -> Dict[str, torch.Tensor]: + """Slice gathered max-rank LoRA tensors down to the active session rank.""" + sliced_state_dict: Dict[str, torch.Tensor] = {} + for name, tensor in lora_state_dict.items(): + rank_dim = cls._infer_lora_rank_dim(name, tensor) + if rank_dim is None: + sliced_state_dict[name] = tensor + continue + + if tensor.shape[rank_dim] < active_rank: + raise ValueError( + f"LoRA tensor {name!r} rank dimension {rank_dim} has size {tensor.shape[rank_dim]}, " + f"which is smaller than requested active_rank={active_rank}." + ) + if tensor.shape[rank_dim] == active_rank: + sliced_state_dict[name] = tensor + continue + + sliced_state_dict[name] = tensor.narrow(rank_dim, 0, active_rank).contiguous() + return sliced_state_dict + + def _save_lora_weights(self, save_path: str, model_id: str, *, preserve_lora_dtype: bool = False) -> None: """ Core LoRA saving logic: activate adapter, gather weights, write PEFT checkpoint. @@ -148,27 +294,58 @@ 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) - # 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: + # Use fast adapter-manager path when available (avoids FSDP unshard). + # Skip the fast path for MoE LoRA: adapter-manager params are rank-local + # under EP, so exporting them directly would drop non-local experts. + configured_target_modules = self.lora_config.get("lora_target_modules") + lora_target_modules = getattr(self, "lora_target_modules", None) or configured_target_modules or [] + has_moe_targets = any(module in lora_target_modules for module in ("gate_proj", "up_proj", "down_proj")) + has_stacked_moe_lora_params = any( + "_lora_" in name and ".experts." in name for name, _param in self.model.named_parameters() + ) + is_moe_lora = bool(self.lora_config.get("moe_hybrid_shared_lora", False)) or ( + has_moe_targets and (has_stacked_moe_lora_params or configured_target_modules is not None) + ) + if self._adapter_manager is not None and model_id in self._adapter_manager.adapters and not is_moe_lora: 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) + logger.info( + f"Rank {self.rank}: Using collective (EP+FSDP2-aware) LoRA save path (is_moe_lora={is_moe_lora})" + ) lora_state_dict = get_lora_state_dict(self.model) # Only rank 0 writes files if self.rank == 0: target_modules, lora_alpha = self._get_lora_save_config() + adapter_session_spec = None + adapter_rank = self.lora_config.get("lora_rank", 32) + lora_export_format = self.lora_config.get("lora_export_format", "peft") + if self._adapter_manager is not None and model_id in self._adapter_manager.adapters: + adapter_session_spec = self._adapter_manager.get_adapter_session_spec(model_id) + adapter_rank = adapter_session_spec["lora_config"]["lora_rank"] + lora_alpha = adapter_session_spec["lora_config"]["lora_alpha"] + # Persist the actual live adapter structure, not the broader + # requested target set from config. Some models expose fused + # projections, so the injected/exported tensors can be a strict + # subset of the requested names. + target_modules = None + lora_state_dict = self._slice_lora_state_dict_to_rank(lora_state_dict, int(adapter_rank)) save_lora_checkpoint( model=self.model, save_path=save_path, base_model_name=self.model_config.get("model_path"), target_modules=target_modules, - r=self.lora_config.get("lora_rank", 32), + r=adapter_rank, lora_alpha=lora_alpha, moe_hybrid_shared_lora=self.lora_config.get("moe_hybrid_shared_lora", False), lora_state_dict=lora_state_dict, + lora_export_format=lora_export_format, + preserve_lora_dtype=preserve_lora_dtype, ) + if adapter_session_spec is not None: + write_session_spec(save_path, adapter_session_spec) # Cleanup del lora_state_dict @@ -206,35 +383,27 @@ def save_adapter_state( if path is None: path = os.path.join(self._adapter_manager.checkpoint_dir, model_id) - # Save LoRA weights (collective operation) - self._save_lora_weights(path, model_id) + local_error = None + try: + # Save LoRA weights (collective operation) + self._save_lora_weights(path, model_id, preserve_lora_dtype=True) - # Only rank 0 writes optimizer and metadata - if self.rank == 0: - # Save optimizer state (adapter-specific) - if save_optimizer: - optimizer_path = os.path.join(path, "optimizer.pt") - torch.save(adapter_state.optimizer.state_dict(), optimizer_path) - - # Save metadata (adapter-specific) - metadata = { - "model_id": model_id, - "global_step": adapter_state.global_step, - "global_forward_backward_step": adapter_state.global_forward_backward_step, - "lr": adapter_state.lr, - "timestamp": time.time(), - "save_optimizer": save_optimizer, - } - metadata_path = os.path.join(path, "metadata.json") - with open(metadata_path, "w") as f: - json.dump(metadata, f, indent=2) + if self.rank == 0: + self._write_adapter_training_artifacts(path, model_id, adapter_state, save_optimizer) + logger.info( + f"Saved adapter state for model_id={model_id} to {path} " + f"(step={adapter_state.global_step}, save_optimizer={save_optimizer})" + ) + except Exception as e: + logger.error(f"Failed to save adapter state for model_id={model_id}: {e}", exc_info=True) + local_error = str(e) - logger.info( - f"Saved adapter state for model_id={model_id} to {path} " - f"(step={adapter_state.global_step}, save_optimizer={save_optimizer})" - ) + synced_error = self._sync_collective_error(local_error) + if synced_error: + raise RuntimeError(f"Adapter state save failed: {synced_error}") - dist.barrier() + if dist.is_available() and dist.is_initialized() and dist.get_world_size() > 1: + dist.barrier() return { "path": path, @@ -550,15 +719,13 @@ 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 = dict(self.model.named_parameters()) # Compute tensor sizes and shard assignments (all ranks compute same assignment) tensor_infos = [] # [(name, estimated_size, is_dtensor), ...] for name, tensor in state_dict_meta.items(): if isinstance(tensor, DTensor): - # For DTensor, compute full size from local shape and mesh - local_shape = tensor.to_local().shape - # Get the sharding spec to compute full shape + # DTensor.shape returns the full logical shape. full_shape = tensor.shape # DTensor.shape returns full logical shape numel = 1 for dim in full_shape: @@ -877,10 +1044,20 @@ def save_lora_only(self, lora_path: str, model_id: str = "default") -> Dict[str, start_time = time.time() - # Save LoRA weights (collective operation) - self._save_lora_weights(lora_path, model_id) + local_error = None + try: + # Save LoRA weights (collective operation) + self._save_lora_weights(lora_path, model_id) + except Exception as e: + logger.error(f"Failed to save LoRA-only checkpoint for model_id={model_id}: {e}", exc_info=True) + local_error = str(e) - dist.barrier() + synced_error = self._sync_collective_error(local_error) + if synced_error: + raise RuntimeError(f"LoRA-only save failed: {synced_error}") + + if dist.is_available() and dist.is_initialized() and dist.get_world_size() > 1: + dist.barrier() # Get step from adapter manager if available if self._adapter_manager is not None: @@ -930,7 +1107,7 @@ def load_state( # For non-LoRA or single-adapter mode, use original DCP approach start_time = time.time() - state = {"model": self.model, "optimizer": self.optimizer if load_optimizer else None, "extra_state": {}} + state = self._build_dcp_load_state(checkpoint_path, load_optimizer=load_optimizer) self.Checkpointer.load(checkpoint_path, state) @@ -939,12 +1116,13 @@ def load_state( self.global_forward_backward_step = state["extra_state"].get("global_forward_backward_step", 0) torch.set_rng_state(state["extra_state"].get("torch_rng_state", torch.get_rng_state())) - dist.barrier() + if dist.is_available() and dist.is_initialized(): + dist.barrier() result = { "checkpoint_path": checkpoint_path, "step": self.global_step, - "load_optimizer": load_optimizer, + "load_optimizer": "optimizer" in state, "load_time": time.time() - start_time, "success": True, } diff --git a/src/xorl/server/runner/model_runner.py b/src/xorl/server/runner/model_runner.py index dec1ef9c..e5b28236 100644 --- a/src/xorl/server/runner/model_runner.py +++ b/src/xorl/server/runner/model_runner.py @@ -15,43 +15,58 @@ import gc import logging +import math import os +import shutil import time +from copy import deepcopy 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 safetensors.torch import save_file +from transformers import AutoTokenizer, PretrainedConfig from xorl.checkpoint import build_checkpointer from xorl.data.constants import IGNORE_INDEX +from xorl.distillation import TeacherActivationCache, TeacherHeadManager 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.sequence_parallel.data import gather_outputs from xorl.lora import LoraLinear from xorl.models.layers.moe.routing_replay import set_replay_stage +from xorl.models.transformers.deepseek_v3.support import deepseek_v3_default_lora_targets from xorl.ops.loss import ( + LossOutput, + OPDLossMetrics, + TokenPartial, causallm_loss_function, importance_sampling_loss_function, + opd_loss_function, policy_loss_function, ) from xorl.optim import build_optimizer from xorl.server.runner.adapters import LoRAAdapterManager from xorl.server.runner.checkpoint import CheckpointManager from xorl.server.runner.utils import MoeMetricsTracker, RoutingReplayHandler, run_self_test, validate_token_ids +from xorl.server.session_spec import build_default_session_spec from xorl.trainers.model_builder import ( build_training_model, resolve_training_model_dtype, ) from xorl.trainers.training_utils import ( clip_gradients, + count_active_microbatches, count_valid_tokens, forward_backward_pp, + get_distsign_grad_scale_factor, + get_effective_grad_clip_value, + make_pp_loss_fn, negotiate_pp_seq_len, pad_micro_batches_for_pp, - pp_loss_fn, + scale_model_gradients, sync_sp_gradients, ) from xorl.trainers.training_utils import ( @@ -86,7 +101,21 @@ def configure_rank0_logging(logger_instance, rank): logger_instance.addFilter(RankFilter(rank)) -def _sp_allreduce_kl_metrics(metrics: Dict[str, Any], sp_group) -> Dict[str, Any]: +def _metric_reduce_op(metric_name: str, metric_ops: Optional[Dict[str, str]] = None) -> str: + if metric_ops and metric_name in metric_ops: + return metric_ops[metric_name] + if metric_name in {"ratio_min", "tis_min"}: + return "min" + if metric_name in {"ratio_max", "tis_max"}: + return "max" + return "mean" + + +def _sp_allreduce_kl_metrics( + metrics: Dict[str, Any], + sp_group, + metric_ops: Optional[Dict[str, str]] = None, +) -> Dict[str, Any]: """ All-reduce KL/ratio metrics across the sequence-parallel (Ulysses) group. @@ -95,49 +124,45 @@ def _sp_allreduce_kl_metrics(metrics: Dict[str, Any], sp_group) -> Dict[str, Any This function aggregates stats across all SP ranks so every rank (especially rank 0 which reports metrics) sees the correct global values. - Mean-type metrics (kl, entropy, ratio_mean, pg_clipfrac) are converted to - (value * local_n) sums, all-reduced with SUM, then divided by total_n. - Min/max metrics use MIN/MAX all-reduce with proper identity elements for - ranks that have no valid tokens. + Mean-type metrics are raw partial sums, so they SUM-reduce and remain + unfinalized. ``valid_tokens`` SUM-reduces alongside them; downstream + accumulation divides mean metrics by the final token count. """ - device = torch.device("cuda") - local_n = metrics.get("_n_valid_kl", 0) - - # --- Mean-type metrics: convert to weighted sums, all-reduce, divide --- - mean_keys = ["kl_sample_train_k3", "entropy_sample", "ratio_mean", "pg_clipfrac"] - sum_tensors = {} - for key in mean_keys: - if key in metrics: - val = float(metrics[key]) * local_n - sum_tensors[key] = torch.tensor(val, dtype=torch.float64, device=device) - - # All-reduce the weighted sums - for t in sum_tensors.values(): - dist.all_reduce(t, op=dist.ReduceOp.SUM, group=sp_group) - - # --- Min/max metrics: use identity elements for empty ranks --- - ratio_min_val = float(metrics.get("ratio_min", 1.0)) if local_n > 0 else float("inf") - ratio_max_val = float(metrics.get("ratio_max", 1.0)) if local_n > 0 else float("-inf") - ratio_min_t = torch.tensor(ratio_min_val, dtype=torch.float64, device=device) - ratio_max_t = torch.tensor(ratio_max_val, dtype=torch.float64, device=device) - dist.all_reduce(ratio_min_t, op=dist.ReduceOp.MIN, group=sp_group) - dist.all_reduce(ratio_max_t, op=dist.ReduceOp.MAX, group=sp_group) - - # --- All-reduce total valid token count --- - n_tensor = torch.tensor(float(local_n), dtype=torch.float64, device=device) - dist.all_reduce(n_tensor, op=dist.ReduceOp.SUM, group=sp_group) - total_n = max(n_tensor.item(), 1.0) + # Backward-compatible argument order for older tests/call sites: + # _sp_allreduce_kl_metrics(metrics, metric_ops, sp_group). + if isinstance(sp_group, dict): + metric_ops, sp_group = sp_group, metric_ops - # --- Update metrics with properly reduced values --- - for key in mean_keys: - if key in sum_tensors: - metrics[key] = sum_tensors[key].item() / total_n + device = torch.device(get_device_type()) + local_n = float(metrics.get("valid_tokens", metrics.get("_n_valid_kl", 0)) or 0) + metrics["valid_tokens"] = local_n - ratio_min_reduced = ratio_min_t.item() - ratio_max_reduced = ratio_max_t.item() - metrics["ratio_min"] = ratio_min_reduced if ratio_min_reduced != float("inf") else 1.0 - metrics["ratio_max"] = ratio_max_reduced if ratio_max_reduced != float("-inf") else 1.0 - metrics["valid_tokens"] = total_n + n_tensor = torch.tensor(local_n, dtype=torch.float64, device=device) + dist.all_reduce(n_tensor, op=dist.ReduceOp.SUM, group=sp_group) + total_n = n_tensor.item() + + for key, value in list(metrics.items()): + if key in {"valid_tokens", "_n_valid_kl"}: + continue + op_name = _metric_reduce_op(key, metric_ops) + if op_name == "min": + local_value = float(value) if local_n > 0 else float("inf") + tensor = torch.tensor(local_value, dtype=torch.float64, device=device) + dist.all_reduce(tensor, op=dist.ReduceOp.MIN, group=sp_group) + reduced = tensor.item() + metrics[key] = tensor if math.isfinite(reduced) else torch.tensor(1.0, dtype=torch.float64, device=device) + elif op_name == "max": + local_value = float(value) if local_n > 0 else float("-inf") + tensor = torch.tensor(local_value, dtype=torch.float64, device=device) + dist.all_reduce(tensor, op=dist.ReduceOp.MAX, group=sp_group) + reduced = tensor.item() + metrics[key] = tensor if math.isfinite(reduced) else torch.tensor(1.0, dtype=torch.float64, device=device) + else: + tensor = torch.as_tensor(value, dtype=torch.float64, device=device).clone() + dist.all_reduce(tensor, op=dist.ReduceOp.SUM, group=sp_group) + metrics[key] = tensor + + metrics["valid_tokens"] = int(total_n) if float(total_n).is_integer() else total_n # Clean up internal key metrics.pop("_n_valid_kl", None) @@ -172,6 +197,32 @@ class ModelRunner: "_original_position_ids", "rollout_logprobs", }, + "opd_loss": { + "labels", + "target_tokens", + "teacher_id", + "teacher_ids", + "teacher_weight", + "teacher_weights", + "teacher_cache_indices", + "teacher_hidden_states", + "_original_position_ids", + }, + "teacher_hidden_cache": { + "labels", + "target_tokens", + "teacher_id", + "teacher_ids", + "teacher_weight", + "teacher_weights", + "teacher_cache_indices", + "teacher_hidden_states", + "_original_position_ids", + "num_samples", + "request_id", + "batch_id", + "_shifted", + }, } def __init__( @@ -202,6 +253,12 @@ def __init__( self.model_config = config.get("model", {}) self.train_config = config.get("train", {}) self.lora_config = config.get("lora", {}) + self._validate_multi_adapter_lora_config() + if self.train_config.get("load_weights_mode") == "skip" and not self.train_config.get("load_checkpoint_path"): + raise ValueError( + "load_weights_mode='skip' skips HF weight loading and requires train.load_checkpoint_path " + "to materialize parameters from a DCP checkpoint." + ) # Cross-entropy mode self.ce_mode = self.train_config.get("ce_mode", "eager") @@ -217,12 +274,23 @@ def __init__( # Deferred gradient normalization: accumulate raw valid token counts # across forward_backward calls, normalize once at optim_step. self._accumulated_valid_tokens: Dict[str, int] = {} + self._accumulated_active_microbatches: Dict[str, int] = {} + self._accumulated_active_voter_total: Dict[str, int] = {} # PP schedule cache: keyed by (n_microbatches, seq_len) to avoid rebuilding on every call. self._pp_schedule_cache: Dict[tuple, Any] = {} + # OPD teacher resource caches are initialized lazily from loss_fn_params. + self._opd_head_manager: Optional[TeacherHeadManager] = None + self._opd_head_config: Optional[Any] = None + self._opd_hidden_cache: Optional[TeacherActivationCache] = None + self._opd_hidden_config: Optional[Any] = None + # Multi-adapter support (initialized later if LoRA is enabled) self._adapter_manager: Optional[LoRAAdapterManager] = None + self._lora_session_specs: Dict[str, Dict[str, Any]] = {} + self._default_lora_session_spec: Optional[Dict[str, Any]] = None + self._checkpoint_mgr: Optional[CheckpointManager] = None # Single-tenant session tracking (for full-weights training mode) # When LoRA is disabled, only one training session is allowed at a time @@ -237,7 +305,9 @@ def __init__( # Device setup get_torch_device().set_device(f"{get_device_type()}:{local_rank}") - helper.set_seed(self.train_config.get("seed", 42), False) + seed = self.train_config.get("seed", 42) + enable_full_determinism = self.train_config.get("enable_full_determinism", False) + helper.set_seed(seed, False) # Disable TF32 and BF16 reduced-precision accumulation for # consistent numerics across parallelism strategies. @@ -250,6 +320,12 @@ def __init__( self._initialize_model() self._initialize_optimizer() self._initialize_checkpointer() + self._checkpoint_mgr = self._build_checkpoint_manager() + self._load_initial_checkpoint() + if enable_full_determinism: + # Enabling deterministic algorithms before Kimi DCP/meta materialization + # makes startup pathologically slow; training and adapter init happen below. + helper.set_seed(seed, True) self._initialize_contexts() # Initialize multi-adapter manager if LoRA is enabled @@ -259,15 +335,29 @@ def __init__( device = torch.device(f"{get_device_type()}:{self.local_rank}") # Only rank 0 should save on eviction to avoid multi-rank file conflicts self._adapter_manager = LoRAAdapterManager( - self.model, device, auto_save_on_eviction=(self.rank == 0), lora_config=self.lora_config + self.model, + device, + checkpoint_dir=self._get_adapter_checkpoint_dir(), + auto_save_on_eviction=(self.rank == 0), + lora_config=self.lora_config, + optimizer_type=self.train_config.get("optimizer", "adamw"), + optimizer_dtype=self.train_config.get("optimizer_dtype", "bf16"), + optimizer_kwargs=self._get_optimizer_kwargs(), + weight_decay=self.train_config.get("weight_decay", 0.01), ) - # Register the "default" adapter with the initial weights and lr - self._adapter_manager.register_adapter( + self._default_lora_session_spec = build_default_session_spec( + base_model=self.model_config.get("model_name") or self.model_config.get("model_path"), + train_config=self.train_config, + lora_config=self.lora_config, + ) + self.register_session( model_id="default", - lr=self.train_config.get("lr", 1e-5), - initialize_fresh=False, # Use current weights as the default adapter + session_spec=self._default_lora_session_spec, + materialize=True, + initialize_fresh=False, ) self._adapter_manager.current_adapter_id = "default" + self._checkpoint_mgr._adapter_manager = self._adapter_manager logger.info("Multi-adapter manager initialized with default adapter") # Initialize tokenizer for sampling (only on rank 0) @@ -280,17 +370,6 @@ def __init__( # Initialize extracted modules self._routing_handler = RoutingReplayHandler(self.model) - self._checkpoint_mgr = CheckpointManager( - model=self.model, - optimizer=self.optimizer, - checkpointer=self.Checkpointer, - lora_config=self.lora_config, - model_config=self.model_config, - train_config=self.train_config, - rank=self.rank, - local_rank=self.local_rank, - 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) @@ -316,6 +395,81 @@ def adapter_manager(self): """Public access to the adapter manager for multi-tenancy LoRA.""" return self._adapter_manager + @property + def lora_session_specs(self) -> Dict[str, Dict[str, Any]]: + """Return the registered LoRA session specs.""" + return self._lora_session_specs + + def get_lora_session_spec(self, model_id: str) -> Dict[str, Any]: + """Get the normalized LoRA session spec for a model_id.""" + if model_id not in self._lora_session_specs: + raise KeyError(f"LoRA session spec not registered for model_id={model_id}") + return deepcopy(self._lora_session_specs[model_id]) + + def _sync_registered_lora_session_spec(self, model_id: str) -> None: + """Refresh the worker session registry from the live adapter state.""" + if self._adapter_manager is None or not self._adapter_manager.has_adapter(model_id): + return + self._lora_session_specs[model_id] = self._adapter_manager.get_adapter_session_spec(model_id) + + def register_session( + self, + model_id: str, + session_spec: Dict[str, Any], + *, + materialize: bool = False, + initialize_fresh: bool = True, + ) -> Dict[str, Any]: + """Register a normalized session runtime spec on this worker.""" + if not self.lora_enabled: + # Full-weight mode remains effectively single-tenant; keep the API + # tolerant of create_model but don't install heterogeneous runtime state. + self._validate_single_tenant(model_id) + return { + "model_id": model_id, + "registered": True, + "materialized": False, + "message": "Full-weight mode ignores per-session LoRA runtime specs.", + } + + existing_spec = self._lora_session_specs.get(model_id) + if existing_spec is not None and existing_spec != session_spec: + raise ValueError( + f"Session '{model_id}' is already registered with a different runtime spec. " + f"existing={existing_spec!r}, requested={session_spec!r}" + ) + + self._lora_session_specs[model_id] = deepcopy(session_spec) + + materialized = False + if materialize and self._adapter_manager is not None and not self._adapter_manager.has_adapter(model_id): + self._adapter_manager.register_adapter( + model_id=model_id, + session_spec=session_spec, + initialize_fresh=initialize_fresh, + ) + materialized = True + + return { + "model_id": model_id, + "registered": True, + "materialized": materialized + or (self._adapter_manager is not None and self._adapter_manager.has_adapter(model_id)), + "session_spec": deepcopy(self._lora_session_specs[model_id]), + } + + def ensure_lora_adapter(self, model_id: str, *, initialize_fresh: bool = True) -> None: + """Materialize a registered LoRA session into the resident adapter registry.""" + if self._adapter_manager is None: + raise RuntimeError("Cannot materialize LoRA adapter: adapter manager not initialized") + session_spec = self.get_lora_session_spec(model_id) + if not self._adapter_manager.has_adapter(model_id): + self._adapter_manager.register_adapter( + model_id=model_id, + session_spec=session_spec, + initialize_fresh=initialize_fresh, + ) + def _check_not_sleeping(self, operation: str) -> None: """Raise if the model is in sleep mode (CPU-offloaded).""" if self.is_sleeping: @@ -348,13 +502,10 @@ def _validate_single_tenant(self, model_id: str) -> None: def kill_session(self, model_id: str, save_checkpoint: bool = True) -> Dict[str, Any]: """ - Kill the active full-weights training session. - - This allows a new session to be started. For LoRA mode, this is a no-op - since multi-tenancy is supported. + Kill an active training session. Args: - model_id: The session to kill (must match active session). + model_id: The session to kill. save_checkpoint: Whether to save a checkpoint before killing. Returns: @@ -364,10 +515,50 @@ def kill_session(self, model_id: str, save_checkpoint: bool = True) -> Dict[str, ValueError: If model_id doesn't match the active session. """ if self.lora_enabled: + if model_id == "default": + return { + "success": True, + "message": "Default LoRA session is reserved and was not removed.", + "checkpoint_path": None, + } + + if model_id not in self._lora_session_specs and ( + self._adapter_manager is None or not self._adapter_manager.has_adapter(model_id) + ): + return { + "success": True, + "message": f"No active LoRA session '{model_id}' to kill.", + "checkpoint_path": None, + } + + checkpoint_path = None + if save_checkpoint and self._adapter_manager is not None: + if self._adapter_manager.has_adapter(model_id): + checkpoint_path = os.path.join( + self.train_config.get("output_dir", "outputs"), + "weights", + model_id, + f"session_{model_id}_final", + ) + self.save_state(checkpoint_path, save_optimizer=True, model_id=model_id) + else: + evicted_path = os.path.join(self._get_adapter_checkpoint_dir(), "evicted", model_id) + if not os.path.exists(evicted_path): + raise FileNotFoundError( + f"Cannot kill LoRA session '{model_id}' with save_checkpoint=True: " + f"adapter is not resident and no evicted checkpoint exists at {evicted_path}." + ) + checkpoint_path = self._promote_evicted_adapter_checkpoint(model_id, evicted_path) + + self._accumulated_valid_tokens.pop(model_id, None) + if self._adapter_manager is not None and self._adapter_manager.has_adapter(model_id): + self._adapter_manager.remove_adapter(model_id) + self._lora_session_specs.pop(model_id, None) + return { "success": True, - "message": "LoRA mode supports multi-tenancy, no session to kill.", - "checkpoint_path": None, + "message": f"LoRA session '{model_id}' killed successfully.", + "checkpoint_path": checkpoint_path, } if self._active_session_id is None: @@ -480,13 +671,14 @@ def _initialize_model(self): 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), + record_routing_weights=self.model_config.get("record_routing_weights", True), 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), init_device=self.train_config.get("init_device", "cpu"), merge_qkv=self.model_config.get("merge_qkv", True), enable_lora=lora_enabled, - lora_rank=self.lora_config.get("lora_rank", 32), + lora_rank=self.lora_config.get("max_lora_rank", self.lora_config.get("lora_rank", 32)), lora_alpha=self.lora_config.get("lora_alpha", 16), lora_target_modules=target_modules, moe_hybrid_shared_lora=self.lora_config.get("moe_hybrid_shared_lora", False), @@ -501,8 +693,9 @@ def _initialize_model(self): basic_modules=self.model_config.get("basic_modules", []), enable_reentrant=self.train_config.get("enable_reentrant", False), enable_forward_prefetch=self.train_config.get("enable_forward_prefetch", True), - load_weights_mode=self.train_config.get("load_weights_mode", "broadcast"), + load_weights_mode=self.train_config.get("load_weights_mode", "grouped"), reshard_after_forward=self.train_config.get("reshard_after_forward"), + moe_grad_reduce_mode=self.train_config.get("moe_grad_reduce_mode", "reduce_scatter"), pp_schedule=pp_schedule_name, freeze_router=self.train_config.get("freeze_router", False), router_fp32=self.model_config.get("router_fp32", True), @@ -511,6 +704,7 @@ def _initialize_model(self): 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), + flash_attention_deterministic=self.model_config.get("flash_attention_deterministic", False), ) self.model = result.model @@ -545,41 +739,125 @@ def _resolve_lora_target_modules(self) -> List[str]: if explicit_target_modules is not None: return explicit_target_modules + config_path = self.model_config.get("config_path") or self.model_config.get("model_path") + model_type = None + if config_path: + try: + config_dict, _ = PretrainedConfig.get_config_dict(config_path) + model_type = config_dict.get("model_type") + if model_type == "kimi_k25": + model_type = config_dict.get("text_config", {}).get("model_type", model_type) + except Exception: + model_type = None + # Legacy Tinker-style: build from boolean flags - target_modules = [] - if self.lora_config.get("train_attn", True): - target_modules.extend(["q_proj", "k_proj", "v_proj", "o_proj"]) - if self.lora_config.get("train_mlp", True): - target_modules.extend(["gate_proj", "up_proj", "down_proj"]) - if self.lora_config.get("train_unembed", True): - target_modules.append("lm_head") + train_attn = self.lora_config.get("train_attn", True) + train_mlp = self.lora_config.get("train_mlp", True) + train_unembed = self.lora_config.get("train_unembed", True) + + if model_type in {"deepseek_v3", "kimi_k2", "kimi_k25"}: + target_modules = deepseek_v3_default_lora_targets( + train_attn=train_attn, + train_mlp=train_mlp, + train_unembed=train_unembed, + ) + else: + target_modules = [] + if train_attn: + target_modules.extend(["q_proj", "k_proj", "v_proj", "o_proj"]) + if train_mlp: + target_modules.extend(["gate_proj", "up_proj", "down_proj"]) + if train_unembed: + target_modules.append("lm_head") if not target_modules: raise ValueError("At least one of train_mlp, train_attn, or train_unembed must be True") return target_modules + def _validate_multi_adapter_lora_config(self) -> None: + """Reject LoRA features that are not supported by the multi-adapter server path.""" + if self.lora_config.get("enable_lora", False) and self.lora_config.get("merge_lora_interval", 0) > 0: + raise ValueError("merge_lora_interval is not supported with multi-adapter LoRA server training") + if self.lora_config.get("enable_lora", False) and self.train_config.get("pipeline_parallel_size", 1) > 1: + raise ValueError( + "pipeline_parallel_size > 1 is not supported with multi-adapter LoRA server training. " + "Adapter coordination currently assumes identical local LoRA layouts on every rank." + ) + max_lora_rank = self.lora_config.get("max_lora_rank", self.lora_config.get("lora_rank", 32)) + default_rank = self.lora_config.get("lora_rank", 32) + if max_lora_rank < default_rank: + raise ValueError( + f"max_lora_rank ({max_lora_rank}) must be >= lora_rank ({default_rank}) for multi-adapter LoRA" + ) + + def _get_optimizer_kwargs(self) -> Dict[str, Any]: + """Collect optimizer kwargs from the server train config.""" + explicit_kwargs = self.train_config.get("optimizer_kwargs") + if explicit_kwargs is not None: + return deepcopy(explicit_kwargs) + + optimizer_type = self.train_config.get("optimizer", "adamw") + kwargs: Dict[str, Any] = {} + if optimizer_type == "muon": + for key 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", + ): + if key in self.train_config: + kwargs[key] = self.train_config[key] + + if self.train_config.get("optimizer_dtype") == "bf16": + kwargs["muon_momentum_dtype"] = torch.bfloat16 + + grad_dtype = self.train_config.get("muon_grad_dtype") + if grad_dtype == "bf16": + kwargs["muon_grad_dtype"] = torch.bfloat16 + elif grad_dtype == "fp32": + kwargs["muon_grad_dtype"] = torch.float32 + + update_dtype = self.train_config.get("muon_update_dtype") + if update_dtype == "bf16": + kwargs["muon_update_dtype"] = torch.bfloat16 + elif update_dtype == "fp32": + kwargs["muon_update_dtype"] = torch.float32 + + if self.train_config.get("muon_force_momentum_path", False): + kwargs["muon_force_momentum_path"] = True + + return kwargs + + def _get_adapter_checkpoint_dir(self) -> str: + """Return the shared adapter checkpoint directory under the server output dir.""" + return os.path.join(self.train_config.get("output_dir", "outputs"), "adapters") + + def _promote_evicted_adapter_checkpoint(self, model_id: str, evicted_path: str) -> str: + """Copy an evicted adapter checkpoint into the public weights namespace.""" + checkpoint_path = os.path.join( + self.train_config.get("output_dir", "outputs"), + "weights", + model_id, + f"session_{model_id}_final", + ) + if self.rank == 0: + os.makedirs(os.path.dirname(checkpoint_path), exist_ok=True) + if os.path.exists(checkpoint_path): + shutil.rmtree(checkpoint_path) + shutil.copytree(evicted_path, checkpoint_path) + return checkpoint_path + def _initialize_optimizer(self): """Initialize the optimizer.""" optimizer_type = self.train_config.get("optimizer", "adamw") - optimizer_kwargs = None - 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", - "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._use_distsignsgd = optimizer_type == "distsignsgd" + if self._use_distsignsgd and self.lora_config.get("enable_lora", False): + raise NotImplementedError("DistSignSGD does not yet support server LoRA adapter-manager training.") + optimizer_kwargs = self._get_optimizer_kwargs() self.optimizer = build_optimizer( self.model, lr=self.train_config.get("lr", 1e-5), @@ -587,7 +865,8 @@ def _initialize_optimizer(self): fused=True, optimizer_type=optimizer_type, optimizer_dtype=self.train_config.get("optimizer_dtype", "bf16"), - optimizer_kwargs=optimizer_kwargs, + optimizer_kwargs=optimizer_kwargs or None, + cautious_weight_decay=self.train_config.get("cautious_weight_decay", False), ) # Register optimizer pre-hook if available @@ -614,16 +893,37 @@ def _initialize_contexts(self): self.train_config.get("activation_gpu_limit", None), ) - def register_lora_adapter(self, model_id: str, lr: float) -> Dict[str, Any]: - """ - Register a new LoRA adapter for a training run. + def _build_checkpoint_manager(self, adapter_manager=None) -> CheckpointManager: + return CheckpointManager( + model=self.model, + optimizer=self.optimizer, + checkpointer=self.Checkpointer, + lora_config=self.lora_config, + model_config=self.model_config, + train_config=self.train_config, + rank=self.rank, + local_rank=self.local_rank, + adapter_manager=adapter_manager, + ) + + def _load_initial_checkpoint(self) -> None: + checkpoint_path = self.train_config.get("load_checkpoint_path") + if not checkpoint_path: + return + if self._checkpoint_mgr is None: + raise RuntimeError("Checkpoint manager must be initialized before loading an initial checkpoint.") + + logger.info("Loading initial checkpoint from %s", checkpoint_path) + self._checkpoint_mgr.load_state(checkpoint_path, load_optimizer=True) + self._sync_from_checkpoint_state() - Creates fresh LoRA weights and optimizer state for the given model_id. - If the model_id already has an adapter, it will be replaced. + def register_lora_adapter(self, model_id: str, lr: Optional[float]) -> Dict[str, Any]: + """ + Materialize a registered LoRA session into the adapter manager. Args: model_id: Unique identifier for this training run - lr: Learning rate for this adapter + lr: Optional learning rate override used for legacy call sites. Returns: Dictionary with registration info @@ -634,11 +934,22 @@ def register_lora_adapter(self, model_id: str, lr: float) -> Dict[str, Any]: if self._adapter_manager is None: raise RuntimeError("Cannot register adapter: LoRA is not enabled or adapter manager not initialized") - self._adapter_manager.register_adapter( - model_id=model_id, - lr=lr, - initialize_fresh=True, - ) + if model_id not in self._lora_session_specs: + if model_id == "default" and self._default_lora_session_spec is not None: + self._lora_session_specs["default"] = deepcopy(self._default_lora_session_spec) + else: + raise KeyError(f"LoRA session '{model_id}' is not registered. Call create_model first.") + + if not self._adapter_manager.has_adapter(model_id): + session_spec = self.get_lora_session_spec(model_id) + if lr is not None: + session_spec["optimizer_config"]["learning_rate"] = lr + self._adapter_manager.register_adapter( + model_id=model_id, + session_spec=session_spec, + initialize_fresh=True, + ) + self._sync_registered_lora_session_spec(model_id) return { "model_id": model_id, @@ -647,7 +958,7 @@ def register_lora_adapter(self, model_id: str, lr: float) -> Dict[str, Any]: "total_adapters": len(self._adapter_manager.list_adapters()), } - def register_adapter(self, model_id: str, lr: float) -> Dict[str, Any]: + def register_adapter(self, model_id: str, lr: Optional[float]) -> Dict[str, Any]: """Alias for register_lora_adapter for API consistency.""" return self.register_lora_adapter(model_id=model_id, lr=lr) @@ -663,12 +974,12 @@ def _get_effective_lm_head_weight(self): return lm_head.weight def _collect_per_token_outputs(self, per_token_tensors, micro_batch, accumulators): - """Gather per-token outputs across Ulysses SP group and append to accumulators.""" + """Gather per-token outputs across the unified SP group and append to accumulators.""" ps = get_parallel_state() if ps.cp_enabled: - ulysses_group = ps.ulysses_group - cp_size = dist.get_world_size(ulysses_group) + sp_group = ps.sp_group + cp_size = ps.cp_size original_position_ids = micro_batch.get("_original_position_ids") if original_position_ids is not None: @@ -687,7 +998,7 @@ def _collect_per_token_outputs(self, per_token_tensors, micro_batch, accumulator padding_dim=-1, unpad_dim_size=original_seq_len, scale_grad=False, - group=ulysses_group, + group=sp_group, ) if position_ids is not None: @@ -711,20 +1022,522 @@ def _collect_per_token_outputs(self, per_token_tensors, micro_batch, accumulator accumulators["losses"].append(gathered["loss"].cpu()) @staticmethod - def _accumulate_is_metrics(accumulated, new_metrics): - """Accumulate importance sampling metrics across micro-batches.""" + def _teacher_cache_dtype(dtype_name: str) -> torch.dtype: + mapping = { + "bf16": torch.bfloat16, + "bfloat16": torch.bfloat16, + "fp16": torch.float16, + "float16": torch.float16, + "fp32": torch.float32, + "float32": torch.float32, + } + key = str(dtype_name).lower() + if key not in mapping: + raise ValueError(f"Unsupported teacher_hidden_cache_dtype={dtype_name!r}") + return mapping[key] + + @staticmethod + def _position_spans(position_ids: torch.Tensor, num_samples: int) -> List[tuple[int, int]]: + pos = position_ids.reshape(-1).to(device="cpu", dtype=torch.long) + if pos.numel() == 0: + return [] + starts = [0] + for i in range(1, pos.numel()): + if pos[i].item() <= pos[i - 1].item(): + starts.append(i) + starts.append(pos.numel()) + return [(starts[i], starts[i + 1]) for i in range(min(num_samples, len(starts) - 1))] + + @staticmethod + def _valid_row_length(labels: Optional[torch.Tensor], row: int, fallback: int) -> int: + if labels is None or labels.numel() == 0: + return fallback + row_labels = labels[row].reshape(-1) + valid = row_labels != IGNORE_INDEX + if valid.any(): + return int(valid.nonzero(as_tuple=True)[0][-1].item()) + 1 + return 0 + + @staticmethod + def _teacher_cache_label_key(micro_batch: Dict[str, Any]) -> Optional[str]: + if isinstance(micro_batch.get("labels"), torch.Tensor): + return "labels" + if isinstance(micro_batch.get("target_tokens"), torch.Tensor): + return "target_tokens" + return None + + def _gather_teacher_cache_sequences( + self, + hidden_states: torch.Tensor, + micro_batch: Dict[str, Any], + ps, + ) -> tuple[torch.Tensor, Dict[str, Any]]: + if not getattr(ps, "cp_enabled", False): + return hidden_states, micro_batch + + original_position_ids = micro_batch.get("_original_position_ids") + if isinstance(original_position_ids, torch.Tensor): + original_seq_len = original_position_ids.shape[-1] + else: + original_seq_len = hidden_states.shape[1] * max(int(getattr(ps, "cp_size", 1)), 1) + + sequence_group = getattr(ps, "sp_group", getattr(ps, "ulysses_group", None)) + hidden_states = gather_outputs( + hidden_states, + gather_dim=1, + padding_dim=1, + unpad_dim_size=original_seq_len, + scale_grad=False, + group=sequence_group, + ) + + label_key = self._teacher_cache_label_key(micro_batch) + if label_key is None: + return hidden_states, micro_batch + + labels = micro_batch[label_key] + if labels.shape[-1] == hidden_states.shape[1]: + return hidden_states, micro_batch + + full_labels = gather_outputs( + labels, + gather_dim=-1, + padding_dim=-1, + unpad_dim_size=original_seq_len, + scale_grad=False, + group=sequence_group, + ) + micro_batch = dict(micro_batch) + micro_batch[label_key] = full_labels + return hidden_states, micro_batch + + def _teacher_cache_write_cp_rank(self, ps, write_rank: int) -> int: + if not getattr(ps, "cp_enabled", False): + return -1 + if not (dist.is_available() and dist.is_initialized()): + return int(getattr(ps, "cp_rank", 0)) + + write_cp_rank = torch.tensor( + [int(getattr(ps, "cp_rank", -1)) if self.rank == write_rank else -1], + dtype=torch.long, + device=get_device_type(), + ) + dist.broadcast(write_cp_rank, src=write_rank) + return int(write_cp_rank.item()) + + def _gather_teacher_cache_chunks( + self, + local_chunks: List[torch.Tensor], + *, + write_rank: int, + slice_key: Optional[int] = 0, + ) -> Optional[List[Dict[str, Any]]]: + payload = {"rank": int(self.rank), "slice_key": int(slice_key or 0), "chunks": local_chunks} + if not (dist.is_available() and dist.is_initialized()): + return [payload] if self.rank == write_rank else None + + gathered = [None for _ in range(dist.get_world_size())] if self.rank == write_rank else None + dist.gather_object(payload, gathered, dst=write_rank) + return [item for item in gathered if item is not None] if self.rank == write_rank else None + + def _write_teacher_hidden_cache( + self, + gathered_payloads: List[Dict[str, Any]], + *, + cache_path: str, + cache_key: str, + cache_dtype: torch.dtype, + now, + ) -> tuple[Dict[str, Any], float]: + chunks, cache_indices_by_sample = self._merge_teacher_hidden_cache_payloads(gathered_payloads) + + if not chunks: + raise ValueError("teacher_hidden_cache produced no hidden-state chunks to write") + + t_write = now() + os.makedirs(os.path.dirname(os.path.abspath(cache_path)) or ".", exist_ok=True) + hidden_cache = torch.cat(chunks, dim=0).contiguous() + tmp_path = f"{cache_path}.tmp-rank{self.rank}" + save_file({cache_key: hidden_cache}, tmp_path) + os.replace(tmp_path, cache_path) + write_s = now() - t_write + + return ( + { + "path": cache_path, + "tensor_key": cache_key, + "dtype": str(cache_dtype).removeprefix("torch."), + "num_tokens": int(hidden_cache.shape[0]), + "hidden_size": int(hidden_cache.shape[1]), + "cache_indices_by_sample": cache_indices_by_sample, + }, + write_s, + ) + + def _teacher_hidden_chunks_from_batch( + self, + hidden_states: torch.Tensor, + micro_batch: Dict[str, Any], + ) -> List[torch.Tensor]: + """Split a teacher forward pass into real per-sample hidden-state chunks.""" + hidden_states = hidden_states.detach() + position_ids = micro_batch.get("_original_position_ids", micro_batch.get("position_ids")) + labels = micro_batch.get("labels", micro_batch.get("target_tokens")) + + if hidden_states.ndim != 3: + raise ValueError(f"Expected teacher hidden states [batch, seq, hidden], got {tuple(hidden_states.shape)}") + + # Packed batches concatenate samples into batch row 0 and record the + # number of real samples. Padding is represented as one extra position-id + # segment, so use num_samples to drop it. + num_samples = int(micro_batch.get("num_samples", 0) or 0) + if num_samples > 0: + if position_ids is None: + raise ValueError("Packed teacher_hidden_cache batches require position_ids") + spans = self._position_spans(position_ids, num_samples) + flat_hidden = hidden_states.reshape(-1, hidden_states.shape[-1]) + return [flat_hidden[start:end].contiguous() for start, end in spans if end > start] + + chunks: List[torch.Tensor] = [] + batch_size = hidden_states.shape[0] + for row in range(batch_size): + fallback_len = hidden_states.shape[1] + length = ( + self._valid_row_length(labels, row, fallback_len) if isinstance(labels, torch.Tensor) else fallback_len + ) + if length > 0: + chunks.append(hidden_states[row, :length].contiguous()) + return chunks + + def _teacher_hidden_cache_contributor_key(self, ps) -> Optional[int]: + """Return this rank's logical cache slice key, or None for duplicate shards.""" + if getattr(ps, "cp_enabled", False) and int(getattr(ps, "cp_rank", 0)) != 0: + return None + + if getattr(ps, "ep_enabled", False): + if int(getattr(ps, "ep_rank", 0)) != 0: + return None + ep_mesh = getattr(ps, "ep_fsdp_device_mesh", None) + if ep_mesh is not None: + try: + return int(ep_mesh.get_local_rank("ep_fsdp")) + except Exception as exc: + logger.debug( + "Rank %s: could not read ep_fsdp local rank from EP mesh (%s); falling back to rank arithmetic", + self.rank, + exc, + ) + + ep_size = max(1, int(getattr(ps, "ep_size", 1))) + ep_fsdp_size = max(1, int(getattr(ps, "dp_shard_in_ep_size", 1))) + pp_size = max(1, int(getattr(ps, "pp_size", 1))) + ranks_per_pp_stage = max(1, self.world_size // pp_size) + local_stage_rank = self.rank % ranks_per_pp_stage + return min(local_stage_rank // ep_size, ep_fsdp_size - 1) + + return int(getattr(ps, "dp_rank", 0)) + + @staticmethod + def _merge_teacher_hidden_cache_payloads(payloads: List[Optional[Dict[str, Any]]]): + """Merge per-rank teacher-cache chunks in logical data-slice order.""" + chunks: List[torch.Tensor] = [] + cache_indices_by_sample: List[List[int]] = [] + next_index = 0 + ordered_payloads = sorted( + (payload for payload in payloads if payload), + key=lambda payload: (int(payload["slice_key"]), int(payload["rank"])), + ) + for payload in ordered_payloads: + for chunk in payload["chunks"]: + rows = int(chunk.shape[0]) + cache_indices_by_sample.append(list(range(next_index, next_index + rows))) + next_index += rows + chunks.append(chunk) + return chunks, cache_indices_by_sample + + def _forward_teacher_hidden_cache( + self, + micro_batches: List[Dict[str, Any]], + params: Dict[str, Any], + abort_callback=None, + ) -> Dict[str, Any]: + if self.pp_enabled: + raise NotImplementedError("teacher_hidden_cache does not yet support pipeline parallelism") + + cache_path = ( + params.get("teacher_hidden_cache_path") or params.get("hidden_cache_path") or params.get("output_path") + ) + if not cache_path: + raise ValueError( + "teacher_hidden_cache requires loss_fn_params.teacher_hidden_cache_path " + "(or hidden_cache_path/output_path)" + ) + + cache_key = params.get("teacher_hidden_cache_key", "hidden_states") + cache_dtype = self._teacher_cache_dtype(params.get("teacher_hidden_cache_dtype", "bfloat16")) + write_rank = int(params.get("teacher_hidden_cache_write_rank", 0)) + world_size = dist.get_world_size() if dist.is_available() and dist.is_initialized() else self.world_size + if write_rank < 0 or write_rank >= world_size: + raise ValueError(f"teacher_hidden_cache_write_rank={write_rank} is outside world_size={world_size}") + + profile_sync_cuda = bool(params.get("opd_profile_sync_cuda", False)) + + def _now() -> float: + if profile_sync_cuda and torch.cuda.is_available(): + torch.cuda.synchronize() + return time.perf_counter() + + ps = get_parallel_state() + contributor_key = self._teacher_hidden_cache_contributor_key(ps) + write_cp_rank = self._teacher_cache_write_cp_rank(ps, write_rank) + is_contributor = contributor_key is not None and ( + not getattr(ps, "cp_enabled", False) or int(getattr(ps, "cp_rank", 0)) == write_cp_rank + ) + + forward_compute_s = 0.0 + local_chunks: List[torch.Tensor] = [] + + for micro_batch in micro_batches: + if abort_callback and abort_callback(): + raise RuntimeError("Execution aborted by request") + + micro_batch = { + k: v.to(get_device_type(), non_blocking=True) if isinstance(v, torch.Tensor) else v + for k, v in micro_batch.items() + } + + model_inputs = { + k: v for k, v in micro_batch.items() if k not in self._LOSS_EXCLUDE_KEYS["teacher_hidden_cache"] + } + + t_forward = _now() + with self.model_fwd_context: + outputs = self.model(**model_inputs, use_cache=False, output_hidden_states=False) + hidden_states, micro_batch = self._gather_teacher_cache_sequences( + outputs.last_hidden_state, micro_batch, ps + ) + forward_compute_s += _now() - t_forward + + if is_contributor: + for chunk in self._teacher_hidden_chunks_from_batch(hidden_states, micro_batch): + local_chunks.append(chunk.to(device="cpu", dtype=cache_dtype)) + + del outputs, hidden_states, micro_batch, model_inputs + + gathered_payloads = self._gather_teacher_cache_chunks( + local_chunks, + write_rank=write_rank, + slice_key=contributor_key, + ) + + cache_metadata = None + write_s = 0.0 + if self.rank == write_rank: + cache_metadata, write_s = self._write_teacher_hidden_cache( + gathered_payloads or [], + cache_path=cache_path, + cache_key=cache_key, + cache_dtype=cache_dtype, + now=_now, + ) + + if dist.is_available() and dist.is_initialized(): + metadata_box = [cache_metadata] + dist.broadcast_object_list(metadata_box, src=write_rank) + cache_metadata = metadata_box[0] + + if cache_metadata is None: + raise RuntimeError("teacher_hidden_cache metadata was not produced") + + result = { + "total_loss": 0.0, + "global_valid_tokens": cache_metadata["num_tokens"], + "teacher_hidden_cache": cache_metadata, + "teacher_prefill_tokens": cache_metadata["num_tokens"], + "teacher_prefill_forward_compute_s": forward_compute_s, + "teacher_hidden_cache_write_s": write_s, + } + return result + + @staticmethod + def _metric_accumulator_key(metric_name: str, loss_fn: str) -> tuple[str, str] | None: + """Return output metric key and reduction mode for a loss metric.""" + if loss_fn == "opd_loss": + if metric_name == "valid_tokens": + # The top-level global_valid_tokens field already reports this. + return None + if metric_name == "opd_num_teachers": + return "opd_num_teachers:max", "max" + if metric_name.startswith("opd_profile_") and metric_name.endswith("_ms"): + return metric_name, "sum_max" + if metric_name.startswith("opd_"): + return metric_name, "mean" + return None + + if metric_name == "ratio_min": + return "is_ratio_min", "min" + if metric_name == "ratio_max": + return "is_ratio_max", "max" + return f"is_{metric_name}", "mean" + + @staticmethod + def _accumulate_loss_metrics(accumulated, new_metrics, loss_fn: str, metric_ops=None): + """Accumulate loss-specific metrics across micro-batches. + + RL loss metrics are reducer partials, so mean metrics already carry a + token-sum share. OPD keeps human-readable per-micro-batch means for its + namespaced metrics, so those are weighted by valid-token count here. + """ if not new_metrics: return + metric_ops = metric_ops or {} + new_metrics = dict(new_metrics) new_metrics.pop("_n_valid_kl", None) - n_tokens = new_metrics.get("valid_tokens", 1) + n_tokens = float(new_metrics.get("valid_tokens", 1)) + device = get_device_type() for k, v in new_metrics.items(): + accumulator_key = ModelRunner._metric_accumulator_key(k, loss_fn) + if accumulator_key is None: + continue + output_key, default_op = accumulator_key + op = metric_ops.get(k, default_op) + value = torch.as_tensor(v, dtype=torch.float64, device=device) + + if op in ("min", "max"): + entry = accumulated.get(output_key) + if entry is None: + accumulated[output_key] = {"value": value.clone(), "op": op} + else: + entry["value"] = ( + torch.minimum(entry["value"], value) if op == "min" else torch.maximum(entry["value"], value) + ) + continue + + if op == "sum_max": + entry = accumulated.get(output_key) + if entry is None: + accumulated[output_key] = {"sum": value.clone(), "op": op} + else: + entry["sum"] = entry["sum"] + value + continue + + if loss_fn == "opd_loss": + value_sum = value * n_tokens + count = n_tokens + elif k == "valid_tokens": + value_sum = value + count = 1.0 + else: + value_sum = value + count = n_tokens + + entry = accumulated.get(output_key) + if entry is None: + accumulated[output_key] = {"sum": value_sum.clone(), "count": float(count), "op": "mean"} + else: + entry["sum"] = entry["sum"] + value_sum + entry["count"] += float(count) + + @staticmethod + def _finalize_loss_metrics(accumulated, result, loss_fn: Optional[str] = None): + """All-reduce loss metrics, then add reduced values to result dict.""" + if not accumulated: + return + ps = get_parallel_state() + if loss_fn is None and all(str(k).startswith("opd_") for k in accumulated): + loss_fn = "opd_loss" + + if loss_fn == "opd_loss": + reduce_group = ps.loss_group if ps.loss_parallel_enabled else None + should_reduce = ps.loss_parallel_enabled and dist.is_available() and dist.is_initialized() + else: + reduce_group = ps.dp_group if ps.dp_enabled else None + should_reduce = ps.dp_enabled and dist.is_available() and dist.is_initialized() + + groups: Dict[str, list[str]] = {"mean": [], "min": [], "max": [], "sum_max": []} + for k, entry in accumulated.items(): + groups[entry["op"]].append(k) + + if groups["mean"]: + keys = groups["mean"] + sums = torch.stack([accumulated[k]["sum"] for k in keys]) + counts = torch.tensor( + [accumulated[k]["count"] for k in keys], + dtype=torch.float64, + device=get_device_type(), + ) + if should_reduce: + sums_and_counts = torch.cat([sums, counts]) + dist.all_reduce(sums_and_counts, op=dist.ReduceOp.SUM, group=reduce_group) + sums, counts = sums_and_counts[: len(keys)], sums_and_counts[len(keys) :] + means = (sums / counts.clamp(min=1.0)).tolist() + mask = (counts > 0).tolist() + for i, k in enumerate(keys): + if mask[i]: + result[k] = means[i] + + for op_name, reduce_op in (("min", dist.ReduceOp.MIN), ("max", dist.ReduceOp.MAX)): + if not groups[op_name]: + continue + keys = groups[op_name] + stacked = torch.stack([accumulated[k]["value"] for k in keys]) + if should_reduce: + dist.all_reduce(stacked, op=reduce_op, group=reduce_group) + values = stacked.tolist() + for k, v in zip(keys, values): + result[k] = v if math.isfinite(v) else 1.0 + + if groups["sum_max"]: + keys = groups["sum_max"] + stacked = torch.stack([accumulated[k]["sum"] for k in keys]) + if should_reduce: + dist.all_reduce(stacked, op=dist.ReduceOp.MAX, group=reduce_group) + values = stacked.tolist() + for k, v in zip(keys, values): + result[k] = v if math.isfinite(v) else 0.0 + + @staticmethod + def _metric_to_float(value): + """Convert scalar metric values to Python floats before cross-process serialization.""" + if isinstance(value, torch.Tensor): + if value.numel() != 1: + raise ValueError(f"Expected scalar metric tensor, got shape {tuple(value.shape)}") + return float(value.detach().cpu().item()) + return float(value) + + @staticmethod + def _accumulate_is_metrics(accumulated, new_metrics, metric_ops: Optional[Dict[str, str]] = None): + """Accumulate importance sampling metrics across micro-batches.""" + if not new_metrics: + return + metric_ops = metric_ops or {} + n_tokens = ModelRunner._metric_to_float(new_metrics.get("valid_tokens", 1)) + for k, raw_v in new_metrics.items(): + if k == "_n_valid_kl": + continue + v = ModelRunner._metric_to_float(raw_v) if k not in accumulated: - accumulated[k] = {"sum": 0.0, "count": 0} + op_name = _metric_reduce_op(k, metric_ops) + if op_name == "min": + accumulated[k] = {"op": op_name, "sum": float("inf"), "count": 0} + elif op_name == "max": + accumulated[k] = {"op": op_name, "sum": float("-inf"), "count": 0} + else: + accumulated[k] = {"op": op_name, "sum": 0.0, "count": 0} + op_name = accumulated[k].get("op", _metric_reduce_op(k, metric_ops)) if k == "valid_tokens": accumulated[k]["sum"] += v accumulated[k]["count"] += 1 + elif op_name == "min": + if n_tokens > 0: + accumulated[k]["sum"] = min(accumulated[k]["sum"], v) + accumulated[k]["count"] += 1 + elif op_name == "max": + if n_tokens > 0: + accumulated[k]["sum"] = max(accumulated[k]["sum"], v) + accumulated[k]["count"] += 1 else: - accumulated[k]["sum"] += v * n_tokens + accumulated[k]["sum"] += v accumulated[k]["count"] += n_tokens @staticmethod @@ -733,28 +1546,33 @@ def _finalize_is_metrics(accumulated, result): if not accumulated: return ps = get_parallel_state() + device = torch.device(get_device_type()) if ps.dp_enabled: dp_group = ps.dp_group for k, v in accumulated.items(): - if k == "ratio_min": - t = torch.tensor(v["sum"] if v["count"] > 0 else float("inf"), dtype=torch.float64, device="cuda") + op_name = v.get("op", _metric_reduce_op(k)) + if op_name == "min": + t = torch.tensor(v["sum"] if v["count"] > 0 else float("inf"), dtype=torch.float64, device=device) dist.all_reduce(t, op=dist.ReduceOp.MIN, group=dp_group) - v["sum"] = t.item() if t.item() != float("inf") else v["sum"] - v["count"] = 1 - elif k == "ratio_max": - t = torch.tensor(v["sum"] if v["count"] > 0 else float("-inf"), dtype=torch.float64, device="cuda") + v["sum"] = t.item() + v["count"] = 1 if math.isfinite(v["sum"]) else 0 + elif op_name == "max": + t = torch.tensor(v["sum"] if v["count"] > 0 else float("-inf"), dtype=torch.float64, device=device) dist.all_reduce(t, op=dist.ReduceOp.MAX, group=dp_group) - v["sum"] = t.item() if t.item() != float("-inf") else v["sum"] - v["count"] = 1 + v["sum"] = t.item() + v["count"] = 1 if math.isfinite(v["sum"]) else 0 else: - sum_t = torch.tensor(v["sum"], dtype=torch.float64, device="cuda") - count_t = torch.tensor(float(v["count"]), dtype=torch.float64, device="cuda") + sum_t = torch.tensor(v["sum"], dtype=torch.float64, device=device) + count_t = torch.tensor(float(v["count"]), dtype=torch.float64, device=device) dist.all_reduce(sum_t, op=dist.ReduceOp.SUM, group=dp_group) dist.all_reduce(count_t, op=dist.ReduceOp.SUM, group=dp_group) v["sum"] = sum_t.item() v["count"] = count_t.item() for k, v in accumulated.items(): - if v["count"] > 0: + op_name = v.get("op", _metric_reduce_op(k)) + if op_name in {"min", "max"}: + result[f"is_{k}"] = v["sum"] if v["count"] > 0 and math.isfinite(v["sum"]) else 1.0 + elif v["count"] > 0: result[f"is_{k}"] = v["sum"] / v["count"] def _count_global_valid_tokens(self, micro_batches): @@ -766,12 +1584,278 @@ def _count_global_valid_tokens(self, micro_batches): group = get_parallel_state().fsdp_group if self.pp_enabled else None return count_valid_tokens(micro_batches, group=group) + def _count_active_microbatches(self, micro_batches) -> tuple[int, int]: + """Return ``(active_microbatches, active_voter_total)`` over the DP group.""" + group = get_parallel_state().fsdp_group if self.pp_enabled else None + return count_active_microbatches(micro_batches, group=group) + + @staticmethod + def _opd_param(params: Dict[str, Any], *names: str, default=None): + for name in names: + if name in params: + return params[name] + return default + + def _get_opd_head_manager(self, params: Dict[str, Any]) -> TeacherHeadManager: + teacher_heads = self._opd_param( + params, + "teacher_heads", + "opd_teacher_heads", + default=self.train_config.get("opd_teacher_heads"), + ) + if not teacher_heads: + raise ValueError("opd_loss requires loss_fn_params.teacher_heads (or train.opd_teacher_heads)") + config_key = repr(teacher_heads) + if self._opd_head_manager is None or self._opd_head_config != config_key: + self._opd_head_manager = TeacherHeadManager(teacher_heads) + self._opd_head_config = config_key + return self._opd_head_manager + + def _get_opd_hidden_cache(self, params: Dict[str, Any]) -> TeacherActivationCache: + hidden_caches = self._opd_param( + params, + "teacher_hidden_caches", + "opd_teacher_hidden_caches", + "teacher_hidden_path", + "opd_teacher_hidden_path", + default=self.train_config.get("opd_teacher_hidden_caches") + or self.train_config.get("opd_teacher_hidden_path"), + ) + if not hidden_caches: + raise ValueError( + "opd_loss requires teacher_hidden_states in the batch or " + "loss_fn_params.teacher_hidden_caches/teacher_hidden_path" + ) + config_key = repr(hidden_caches) + if self._opd_hidden_cache is None or self._opd_hidden_config != config_key: + self._opd_hidden_cache = TeacherActivationCache(hidden_caches) + self._opd_hidden_config = config_key + return self._opd_hidden_cache + + def _get_opd_teacher_hidden_states( + self, + micro_batch: Dict[str, Any], + teacher_id: int, + params: Dict[str, Any], + dtype: torch.dtype, + teacher_mask: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + if "teacher_hidden_states" in micro_batch: + return micro_batch["teacher_hidden_states"].to(get_device_type(), dtype=dtype, non_blocking=True) + + cache_indices = micro_batch.get("teacher_cache_indices") + if cache_indices is None: + raise ValueError("opd_loss requires teacher_cache_indices when teacher_hidden_states are not provided") + + if teacher_mask is not None: + mask = teacher_mask.to(device=cache_indices.device) + selected_indices = cache_indices[mask] + if selected_indices.numel() == 0: + raise ValueError(f"No cache indices available for teacher_id={teacher_id}") + cache_indices = cache_indices.masked_fill(~mask, selected_indices.reshape(-1)[0]) + + cache = self._get_opd_hidden_cache(params) + return cache.get(teacher_id, cache_indices, device=get_device_type(), dtype=dtype) + + def _compute_opd_micro_batch_loss( + self, + hidden_states: torch.Tensor, + student_weight: torch.Tensor, + micro_batch: Dict[str, Any], + params: Dict[str, Any], + loss_reducer=None, + student_lm_head=None, + ) -> LossOutput: + if get_parallel_state().tp_enabled: + raise NotImplementedError("opd_loss does not yet support tensor parallelism") + if self.pp_enabled: + # Mirrors the dispatcher-level guard in _run_forward_backward / _run_forward — + # belt-and-suspenders so a future direct caller can't accidentally launch + # OPD under PP without that pathway being thought through. + raise NotImplementedError("opd_loss does not yet support pipeline parallelism") + + labels = micro_batch.get("labels", micro_batch.get("target_tokens")) + if labels is None: + raise ValueError("opd_loss requires labels or target_tokens for its valid-token mask") + + teacher_ids = micro_batch.get("teacher_ids") + if teacher_ids is None: + default_teacher_id = int(params.get("teacher_id", 0)) + teacher_ids = torch.full_like(labels, default_teacher_id) + else: + teacher_ids = teacher_ids.to(labels.device) + + teacher_weights = micro_batch.get("teacher_weights") + if teacher_weights is None and "teacher_weight" in params: + teacher_weights = torch.full(labels.shape, float(params["teacher_weight"]), device=labels.device) + + valid_mask = labels != IGNORE_INDEX + local_valid_tokens = valid_mask.sum() + lm_head_anchor = self._lm_head_forward_anchor(hidden_states, student_lm_head) + if local_valid_tokens.item() == 0: + # fp32 to match opd_loss_function's fp32 normal-return path. With + # ulysses sequence sharding, ranks with no response tokens hit this + # branch while the rank holding the response runs the fp32 KL kernel. + # A dtype mismatch corrupts the cross-rank all_reduce in the loss + # reporter (mixed-dtype byte reinterpretation). + loss = hidden_states.float().sum() * 0.0 + student_weight.float().sum() * 0.0 + lm_head_anchor + return LossOutput(loss=loss, metrics=OPDLossMetrics(valid_tokens=0).to_dict()) + + head_manager = self._get_opd_head_manager(params) + unique_teacher_ids = sorted(int(x) for x in torch.unique(teacher_ids[valid_mask]).tolist()) + kl_backend = params.get("opd_kl_backend", params.get("kl_backend", "torch_compile")) + use_sharded_backend = str(kl_backend).lower() in {"streaming", "tilelang"} + async_prefetch = bool(params.get("opd_async_prefetch", True)) + teacher_head_fp32 = bool(params.get("teacher_lm_head_fp32", True)) + teacher_head_dtype = torch.float32 if teacher_head_fp32 else student_weight.dtype + sharded_head_cpu_cache = bool(params.get("opd_sharded_head_cpu_cache", True)) + sharded_head_device_cache = bool(params.get("opd_sharded_head_device_cache", False)) + profile_timings = bool(params.get("opd_profile_timings", False)) + profile_sync_cuda = bool(params.get("opd_profile_sync_cuda", False)) + + def _profile_now() -> float: + if profile_sync_cuda and torch.cuda.is_available(): + torch.cuda.synchronize() + return time.perf_counter() + + def _profile_elapsed_ms(start: float) -> float: + return (_profile_now() - start) * 1000.0 + + profile_total_start = _profile_now() if profile_timings else 0.0 + profile_prefetch_ms = 0.0 + profile_hidden_fetch_ms = 0.0 + profile_head_prepare_ms = 0.0 + profile_kl_compute_ms = 0.0 + + hidden_cache = None + if async_prefetch: + profile_start = _profile_now() if profile_timings else 0.0 + if not (use_sharded_backend and head_manager.has_sharded_head(unique_teacher_ids[0])): + head_manager.prefetch(unique_teacher_ids[0]) + if "teacher_hidden_states" not in micro_batch: + hidden_cache = self._get_opd_hidden_cache(params) + hidden_cache.prefetch(unique_teacher_ids[0]) + if profile_timings: + profile_prefetch_ms += _profile_elapsed_ms(profile_start) + if loss_reducer is None: + loss_reducer = TokenPartial(scale=torch.tensor(1.0, device=hidden_states.device)) + + # fp32 — matches opd_loss_function's fp32 return; consistent across ranks. + total_loss = hidden_states.float().sum() * 0.0 + student_weight.float().sum() * 0.0 + lm_head_anchor + weighted_kl_metric = 0.0 + kl_sum = 0.0 + teacher_weight_sum = 0.0 + valid_count = int(local_valid_tokens.item()) + + for teacher_index, teacher_id_int in enumerate(unique_teacher_ids): + if async_prefetch and teacher_index + 1 < len(unique_teacher_ids): + profile_start = _profile_now() if profile_timings else 0.0 + next_teacher_id = unique_teacher_ids[teacher_index + 1] + if not (use_sharded_backend and head_manager.has_sharded_head(next_teacher_id)): + head_manager.prefetch(next_teacher_id) + if hidden_cache is not None: + hidden_cache.prefetch(next_teacher_id) + if profile_timings: + profile_prefetch_ms += _profile_elapsed_ms(profile_start) + teacher_mask = valid_mask & (teacher_ids == teacher_id_int) + if not teacher_mask.any(): + continue + + group_labels = labels.masked_fill(~teacher_mask, IGNORE_INDEX) + profile_start = _profile_now() if profile_timings else 0.0 + teacher_hidden_states = self._get_opd_teacher_hidden_states( + micro_batch, + teacher_id_int, + params, + dtype=hidden_states.dtype, + teacher_mask=teacher_mask, + ) + if profile_timings: + profile_hidden_fetch_ms += _profile_elapsed_ms(profile_start) + + profile_start = _profile_now() if profile_timings else 0.0 + if use_sharded_backend and head_manager.has_sharded_head(teacher_id_int): + teacher_head = head_manager.sharded_view( + teacher_id_int, + device=get_device_type(), + dtype=teacher_head_dtype, + cache_cpu=sharded_head_cpu_cache, + cache_device=sharded_head_device_cache, + ) + else: + teacher_head = head_manager.get( + teacher_id_int, + device=get_device_type(), + dtype=teacher_head_dtype, + ) + if profile_timings: + profile_head_prepare_ms += _profile_elapsed_ms(profile_start) + + profile_start = _profile_now() if profile_timings else 0.0 + result = opd_loss_function( + hidden_states=hidden_states, + weight=student_weight, + labels=group_labels, + teacher_hidden_states=teacher_hidden_states, + teacher_lm_head_weight=teacher_head, + teacher_weights=teacher_weights, + ignore_index=IGNORE_INDEX, + num_chunks=params.get("num_chunks", 8), + lm_head_fp32=self.lm_head_fp32, + teacher_lm_head_fp32=teacher_head_fp32, + kl_backend=kl_backend, + vocab_chunk_size=params.get("opd_vocab_chunk_size", params.get("vocab_chunk_size", 32768)), + return_per_token=False, + loss_reducer=loss_reducer, + ) + if profile_timings: + profile_kl_compute_ms += _profile_elapsed_ms(profile_start) + total_loss = total_loss + result.loss + + metrics = result.metrics or {} + group_valid = int(metrics.get("valid_tokens", teacher_mask.sum().item())) + kl_sum += float(metrics.get("opd_kl", 0.0)) * group_valid + teacher_weight_sum += float(metrics.get("opd_teacher_weight_mean", 1.0)) * group_valid + weighted_kl_metric += float(metrics.get("opd_weighted_kl", 0.0)) * group_valid + + metrics = OPDLossMetrics( + valid_tokens=valid_count, + opd_kl=kl_sum / max(valid_count, 1), + opd_weighted_kl=weighted_kl_metric / max(valid_count, 1), + opd_teacher_weight_mean=teacher_weight_sum / max(valid_count, 1), + opd_num_teachers=len(unique_teacher_ids), + ).to_dict() + if profile_timings: + metrics.update( + { + "opd_profile_prefetch_ms": profile_prefetch_ms, + "opd_profile_hidden_fetch_ms": profile_hidden_fetch_ms, + "opd_profile_head_prepare_ms": profile_head_prepare_ms, + "opd_profile_kl_compute_ms": profile_kl_compute_ms, + "opd_profile_total_ms": _profile_elapsed_ms(profile_total_start), + } + ) + + return LossOutput( + loss=total_loss, + metrics=metrics, + ) + # ========================================================================= # Loss computation dispatch # ========================================================================= + @staticmethod + def _lm_head_forward_anchor(hidden_states: torch.Tensor, lm_head) -> torch.Tensor: + """Run lm_head.forward with zero loss contribution for FSDP hook ordering.""" + if lm_head is None or hidden_states.numel() == 0: + return hidden_states.float().sum() * 0.0 + hidden_flat = hidden_states.reshape(-1, hidden_states.shape[-1]) + logits = lm_head(hidden_flat[:1]) + return logits.float().sum() * 0.0 + def _compute_micro_batch_loss(self, micro_batch, loss_fn, loss_fn_params): - """Compute loss for a single micro-batch. Returns (loss, per_token_outputs_dict, is_metrics, model_outputs).""" + """Compute loss for a single micro-batch.""" params = loss_fn_params or {} return_per_token = params.get("return_per_token", True) @@ -782,8 +1866,13 @@ def _compute_micro_batch_loss(self, micro_batch, loss_fn, loss_fn_params): hidden_states = outputs.last_hidden_state effective_weight = self._get_effective_lm_head_weight() + # scale=1 → loss_fns return raw masked sums; normalization deferred to + # optim_step / _finalize_is_metrics. + token_sum_reducer = TokenPartial(scale=torch.tensor(1.0, device=hidden_states.device)) + per_token_outputs = {} is_metrics = None + is_metric_ops = None if loss_fn in ["causallm_loss", "cross_entropy"]: labels = micro_batch.get("labels") @@ -795,7 +1884,7 @@ def _compute_micro_batch_loss(self, micro_batch, loss_fn, loss_fn_params): ce_mode=self.ce_mode, lm_head_fp32=self.lm_head_fp32, ) - loss = _result.loss + local_loss_sum = _result.loss if return_per_token: per_token_outputs["logprobs"] = _result.per_token_logprobs per_token_outputs["loss"] = _result.per_token_loss @@ -815,13 +1904,21 @@ def _compute_micro_batch_loss(self, micro_batch, loss_fn, loss_fn_params): ce_mode=self.ce_mode, compute_kl_stats=compute_kl_stats, lm_head_fp32=self.lm_head_fp32, + metric_reducer=token_sum_reducer, ) - loss = _result.loss + local_loss_sum = _result.loss per_token_outputs["logprobs"] = _result.per_token_logprobs is_metrics = _result.metrics + is_metric_ops = _result.metric_ops + if is_metrics is not None and "valid_tokens" not in is_metrics: + is_metrics["valid_tokens"] = int((target_tokens != IGNORE_INDEX).sum().item()) if compute_kl_stats and get_parallel_state().cp_enabled and is_metrics: - is_metrics = _sp_allreduce_kl_metrics(is_metrics, get_parallel_state().ulysses_group) + is_metrics = _sp_allreduce_kl_metrics( + is_metrics, + get_parallel_state().ulysses_group, + is_metric_ops, + ) # Diagnostic top-k extraction (forward-only feature, rarely used) diagnostic_topk = params.get("diagnostic_topk", 0) @@ -913,18 +2010,36 @@ def _compute_micro_batch_loss(self, micro_batch, loss_fn, loss_fn_params): compute_kl_stats=compute_kl_stats, lm_head_fp32=self.lm_head_fp32, icepop_beta=icepop_beta, + metric_reducer=token_sum_reducer, ) - loss = _result.loss + local_loss_sum = _result.loss per_token_outputs["logprobs"] = _result.per_token_logprobs is_metrics = _result.metrics + is_metric_ops = _result.metric_ops if compute_kl_stats and get_parallel_state().cp_enabled and is_metrics: - is_metrics = _sp_allreduce_kl_metrics(is_metrics, get_parallel_state().ulysses_group) + is_metrics = _sp_allreduce_kl_metrics( + is_metrics, + get_parallel_state().ulysses_group, + is_metric_ops, + ) + + elif loss_fn == "opd_loss": + _result = self._compute_opd_micro_batch_loss( + hidden_states=hidden_states, + student_weight=effective_weight, + micro_batch=micro_batch, + params=params, + loss_reducer=token_sum_reducer, + student_lm_head=getattr(self.model, "lm_head", None), + ) + local_loss_sum = _result.loss + is_metrics = _result.metrics else: raise ValueError(f"Unknown loss_fn: {loss_fn}") - return loss, per_token_outputs, is_metrics, outputs + return local_loss_sum, per_token_outputs, is_metrics, is_metric_ops, outputs # ========================================================================= # Per-sample K3 KL divergence @@ -993,16 +2108,42 @@ def _forward_loop( ): """Core forward (+ optional backward) loop shared between forward and forward_backward.""" params = loss_fn_params or {} + use_distsignsgd = getattr(self, "_use_distsignsgd", False) + + if loss_fn == "teacher_hidden_cache": + if compute_backward: + raise ValueError("teacher_hidden_cache is a forward-only operation") + try: + return self._forward_teacher_hidden_cache(micro_batches, params, abort_callback=abort_callback) + finally: + if r3_enabled: + self._routing_handler.cleanup() # Count valid tokens globally global_valid_tokens = self._count_global_valid_tokens(micro_batches) + if compute_backward and use_distsignsgd: + active_microbatches, active_voter_total = self._count_active_microbatches(micro_batches) + else: + active_microbatches, active_voter_total = 0, 0 if abort_callback and abort_callback(): raise RuntimeError("Execution aborted by request") total_loss = 0.0 - accumulated_is_metrics = {} + accumulated_loss_metrics = {} accumulators = {"logprobs": [], "losses": [], "position_ids": []} + profile_phase_timings = bool(params.get("profile_phase_timings", params.get("opd_profile_timings", False))) + profile_sync_cuda = bool(params.get("opd_profile_sync_cuda", False)) + forward_compute_time = 0.0 + backward_compute_time = 0.0 + + def _profile_phase_now() -> float: + if profile_sync_cuda and torch.cuda.is_available(): + torch.cuda.synchronize() + return time.perf_counter() + + def _profile_phase_elapsed(start: float) -> float: + return _profile_phase_now() - start # Per-sample K3 deferred computation compute_per_sample_k3 = params.get("compute_per_sample_k3", False) @@ -1027,14 +2168,17 @@ def _forward_loop( set_replay_stage("replay_forward") # Forward pass + loss computation + profile_start = _profile_phase_now() if profile_phase_timings else 0.0 with self.model_fwd_context: - loss, per_token_outputs, is_metrics, outputs = self._compute_micro_batch_loss( + local_loss_sum, per_token_outputs, is_metrics, is_metric_ops, outputs = self._compute_micro_batch_loss( micro_batch, loss_fn, params ) + if profile_phase_timings: + forward_compute_time += _profile_phase_elapsed(profile_start) logger.debug( f"Rank {self.rank}: micro_batch {batch_idx}/{len(micro_batches)} " - f"loss={loss.item():.6f}, local_valid_tokens={local_valid_tokens.item()}, " + f"loss_sum={local_loss_sum.item():.6f}, local_valid_tokens={local_valid_tokens.item()}, " f"global_valid_tokens={global_valid_tokens.item()}" ) # Note: loss is always finite even when local_valid_tokens=0, because @@ -1092,7 +2236,6 @@ def _forward_loop( # which is critical for FSDP2 reduce-scatter collectives. if compute_backward: ps = get_parallel_state() - raw_loss = loss * local_valid_tokens.detach().float() if abort_callback and abort_callback(): raise RuntimeError("Execution aborted by request") @@ -1101,27 +2244,32 @@ def _forward_loop( if r3_enabled: set_replay_stage("replay_backward") + profile_start = _profile_phase_now() if profile_phase_timings else 0.0 with self.model_bwd_context: - raw_loss.backward() + local_loss_sum.backward() + if profile_phase_timings: + backward_compute_time += _profile_phase_elapsed(profile_start) # Loss reporting (separately, no grad): compute normalized per-token loss with torch.no_grad(): - loss_report = loss.detach() * local_valid_tokens + # Cast to fp32 before the cross-rank reduction: with ulysses SP, + # ranks holding only IGNORE_INDEX tokens may hit early-return + # paths with a different dtype than the normal-return path. + loss_report = local_loss_sum.detach().float() dist.all_reduce(loss_report, op=dist.ReduceOp.SUM, group=ps.fsdp_group if self.pp_enabled else None) if global_valid_tokens.item() > 0: total_loss += (loss_report / global_valid_tokens).item() else: # Forward-only: accumulate weighted loss if global_valid_tokens.item() > 0: - total_loss += loss.item() * (local_valid_tokens.item() / global_valid_tokens.item()) + total_loss += local_loss_sum.item() / global_valid_tokens.item() - # Accumulate IS metrics - self._accumulate_is_metrics(accumulated_is_metrics, is_metrics) + # Accumulate loss-specific metrics. RL losses keep the historical + # `is_*` prefix; OPD metrics are already namespaced as `opd_*`. + self._accumulate_loss_metrics(accumulated_loss_metrics, is_metrics, loss_fn, is_metric_ops) # Cleanup - del micro_batch, outputs, loss - if compute_backward: - del raw_loss + del micro_batch, outputs, local_loss_sum # Note: gc.collect() + empty_cache() removed from per-step path. # They cost ~250ms + ~50ms per step (profiled on Qwen3-8B 8xH100). @@ -1133,17 +2281,44 @@ def _forward_loop( # CP/SP gradient sync (backward only) if compute_backward: - sync_sp_gradients(self.model, get_parallel_state().sp_grad_sync_group) + sync_sp_gradients( + self.model, + get_parallel_state().sp_grad_sync_group, + skip_dtensor_grads=use_distsignsgd, + ) # Accumulate valid tokens for deferred normalization at optim_step self._accumulated_valid_tokens[model_id] = ( self._accumulated_valid_tokens.get(model_id, 0) + global_valid_tokens.item() ) + if use_distsignsgd: + self._accumulated_active_microbatches[model_id] = ( + self._accumulated_active_microbatches.get(model_id, 0) + active_microbatches + ) + self._accumulated_active_voter_total[model_id] = ( + self._accumulated_active_voter_total.get(model_id, 0) + active_voter_total + ) + + if profile_phase_timings and dist.is_available() and dist.is_initialized(): + phase_times = torch.tensor( + [forward_compute_time, backward_compute_time], + dtype=torch.float64, + device=get_device_type(), + ) + dist.all_reduce(phase_times, op=dist.ReduceOp.MAX) + forward_compute_time = float(phase_times[0].item()) + backward_compute_time = float(phase_times[1].item()) # Build result result = { "total_loss": total_loss, "global_valid_tokens": global_valid_tokens.item(), } + if profile_phase_timings: + result["forward_compute_time"] = forward_compute_time + result["backward_compute_time"] = backward_compute_time + if loss_fn == "opd_loss": + result["opd_profile_forward_compute_s"] = forward_compute_time + result["opd_profile_backward_compute_s"] = backward_compute_time if accumulators["logprobs"]: result["packed_logprobs"] = [t.tolist() for t in accumulators["logprobs"]] @@ -1160,8 +2335,8 @@ def _forward_loop( all_per_sample_k3.extend(per_sample) result["per_sample_k3"] = all_per_sample_k3 - # All-reduce IS metrics across DP and add to result - self._finalize_is_metrics(accumulated_is_metrics, result) + # All-reduce loss-specific metrics across DP and add to result. + self._finalize_loss_metrics(accumulated_loss_metrics, result, loss_fn) synchronize() return result @@ -1189,7 +2364,7 @@ def _get_pp_schedule(self, n_microbatches, seq_len): self._pp_schedule_cache[key] = build_pipeline_schedule( stages=[stage], n_microbatches=n_microbatches, - loss_fn=pp_loss_fn, + loss_fn=make_pp_loss_fn(self.ce_mode), schedule_name=self.train_config.get("pipeline_parallel_schedule", "1F1B"), ) return self._pp_schedule_cache[key] @@ -1283,7 +2458,7 @@ def forward_backward( # Switch to the correct adapter for this model_id if self._adapter_manager is not None: - self._adapter_manager.switch_adapter(model_id, auto_register=True) + self._adapter_manager.switch_adapter(model_id) # Validate token IDs before processing to catch out-of-vocab errors early # This prevents CUDA device-side asserts that can hang the server @@ -1293,6 +2468,10 @@ def forward_backward( # Get return_per_token flag from loss_fn_params (default True for tinker compatibility) params = loss_fn_params or {} + use_distsignsgd = getattr(self, "_use_distsignsgd", False) + + if self.pp_enabled and loss_fn == "opd_loss": + raise NotImplementedError("opd_loss does not yet support pipeline parallelism") # Reference forward pass: compute Xorl's own logprobs to replace SGLang logprobs # This guarantees KL=0 at step 0 since both old and new logprobs come from the same engine @@ -1383,6 +2562,10 @@ def forward_backward( # PP path if self.pp_enabled: global_valid_tokens = self._count_global_valid_tokens(micro_batches) + if use_distsignsgd: + active_microbatches, active_voter_total = self._count_active_microbatches(micro_batches) + else: + active_microbatches, active_voter_total = 0, 0 # Static padding: pad to sample_packing_sequence_len upfront. # With pp_variable_seq_lengths, padding is deferred to _forward_backward_pp. if not self.train_config.get("pp_variable_seq_lengths", False): @@ -1403,6 +2586,13 @@ def forward_backward( } # Accumulate valid tokens for deferred normalization at optim_step self._accumulated_valid_tokens[model_id] = self._accumulated_valid_tokens.get(model_id, 0) + gvt + if use_distsignsgd: + self._accumulated_active_microbatches[model_id] = ( + self._accumulated_active_microbatches.get(model_id, 0) + active_microbatches + ) + self._accumulated_active_voter_total[model_id] = ( + self._accumulated_active_voter_total.get(model_id, 0) + active_voter_total + ) # R3 cleanup for PP path (stage management handled by _pp_forward) if r3_enabled: self._routing_handler.cleanup() @@ -1435,6 +2625,23 @@ def forward_backward( result["forward_backward_time"] = time.time() - start_time result["model_id"] = model_id + if params.get("profile_clear_gradients_after_backward", False): + clear_start = time.perf_counter() + synchronize() + if self.optimizer is not None: + try: + self.optimizer.zero_grad(set_to_none=True) + except TypeError: + self.optimizer.zero_grad() + if self.model is not None: + self.model.zero_grad(set_to_none=True) + self._accumulated_valid_tokens[model_id] = 0 + self._accumulated_active_microbatches[model_id] = 0 + self._accumulated_active_voter_total[model_id] = 0 + gc.collect() + torch.cuda.empty_cache() + result["opd_profile_clear_gradients_ms"] = (time.perf_counter() - clear_start) * 1000.0 + # Increment step counter (use adapter manager if available, else global) if self._adapter_manager is not None: self._adapter_manager.increment_forward_backward_step(model_id) @@ -1473,21 +2680,35 @@ def forward_backward( return result + # FSDP2's pre-forward unshard path preserves parameter version counters. + # torch.inference_mode() disables those counters and breaks forward-only + # teacher_hidden_cache requests, so use no_grad() for eval-style forwards. @torch.no_grad() def forward( self, micro_batches: List[Dict[str, Any]], loss_fn: str = "causallm_loss", loss_fn_params: Optional[Dict[str, Any]] = None, + model_id: str = "default", routed_experts: Optional[List] = None, routed_expert_logits: Optional[List[Any]] = None, ) -> Dict[str, Any]: """Execute forward pass only (no gradient computation).""" self._check_not_sleeping("forward") + + # Validation/eval requests must run against the same tenant adapter as + # training requests for that model_id. + self._validate_single_tenant(model_id) + if self._adapter_manager is not None: + self._adapter_manager.switch_adapter(model_id) + validate_token_ids(micro_batches, self.model.config.vocab_size) start_time = time.time() + if self.pp_enabled and loss_fn == "opd_loss": + raise NotImplementedError("opd_loss does not yet support pipeline parallelism") + r3_enabled = self._routing_handler.setup(micro_batches, routed_experts, routed_expert_logits) result = self._forward_loop( @@ -1496,20 +2717,27 @@ def forward( loss_fn_params, compute_backward=False, r3_enabled=r3_enabled, + model_id=model_id, ) - result["step"] = self.global_forward_backward_step + if self._adapter_manager is not None: + result["step"] = self._adapter_manager.get_adapter_state(model_id).global_forward_backward_step + else: + result["step"] = self.global_forward_backward_step result["forward_time"] = time.time() - start_time + result["model_id"] = model_id logger.info( f"forward loss={result['total_loss']:.4f} " f"tokens={result.get('global_valid_tokens', 'N/A')} " + f"model_id={model_id} " f"time={result['forward_time']:.2f}s" ) logger.debug( f"Rank {self.rank}: forward loss={result['total_loss']:.6f}, " f"global_valid_tokens={result.get('global_valid_tokens', 'N/A')}, " f"n_micro_batches={len(micro_batches)}, loss_fn={loss_fn}, " + f"model_id={model_id}, " f"time={result['forward_time']:.3f}s" ) return result @@ -1560,6 +2788,9 @@ def optim_step( # Pop accumulated valid tokens for this model_id (deferred normalization) accumulated = self._accumulated_valid_tokens.pop(model_id, 0) + accumulated_active_microbatches = getattr(self, "_accumulated_active_microbatches", {}).pop(model_id, 0) + accumulated_active_voter_total = getattr(self, "_accumulated_active_voter_total", {}).pop(model_id, 0) + use_distsignsgd = getattr(self, "_use_distsignsgd", False) # Multi-adapter path: use adapter's own optimizer on adapter's own parameters if self._adapter_manager is not None: @@ -1578,18 +2809,20 @@ def optim_step( clip_value, accumulated_valid_tokens=accumulated, ) + self._sync_registered_lora_session_spec(model_id) current_step = self._adapter_manager.get_global_step(model_id) current_lr = effective_lr # Single-adapter path: use shared optimizer on model parameters else: - # Deferred gradient normalization: scale raw gradients by 1/accumulated_valid_tokens - # Use in-place mul_ to preserve DTensor metadata (FSDP2 grads are DTensors). - if accumulated > 0: - scale = 1.0 / accumulated - for p in self.model.parameters(): - if p.grad is not None: - p.grad.mul_(scale) + if use_distsignsgd: + if accumulated_active_voter_total > 0: + scale_model_gradients( + self.model, + get_distsign_grad_scale_factor(accumulated_active_voter_total), + ) + elif accumulated > 0: + scale_model_gradients(self.model, 1.0 / float(accumulated)) # Determine learning rate if lr is not None: @@ -1598,6 +2831,10 @@ def optim_step( param_group["lr"] = effective_lr ps = get_parallel_state() + clip_value = get_effective_grad_clip_value( + clip_value, + use_distsignsgd=use_distsignsgd, + ) grad_norm = clip_gradients( self.model, @@ -1608,6 +2845,10 @@ def optim_step( # Optimizer step self.optimizer.step() + # Fused/foreach optimizer kernels can still be reading gradients when + # Python reaches zero_grad/empty_cache. Synchronize before releasing + # grad storage to avoid allocator reuse while those kernels are live. + synchronize() try: self.optimizer.zero_grad(set_to_none=True) except TypeError: @@ -1642,6 +2883,8 @@ def optim_step( f"Rank {self.rank}: optim_step step={current_step}, " f"grad_norm={grad_norm:.6f}, lr={current_lr:.2e}, " f"clip={clip_value}, accumulated_valid_tokens={accumulated}, " + f"accumulated_active_microbatches={accumulated_active_microbatches}, " + f"accumulated_active_voter_total={accumulated_active_voter_total}, " f"model_id={model_id}, time={result['optim_step_time']:.3f}s" ) @@ -1651,6 +2894,11 @@ def optim_step( def _maybe_merge_lora(self) -> None: """Periodic LoRA merge at merge_lora_interval.""" + if self._adapter_manager is not None: + merge_interval = self.lora_config.get("merge_lora_interval", 0) + if merge_interval > 0: + raise RuntimeError("merge_lora_interval is not supported with multi-adapter LoRA server training") + return _maybe_merge_lora_util( self.model, enable_lora=self.lora_config.get("enable_lora", False), @@ -1684,6 +2932,7 @@ def save_adapter_state(self, model_id, path=None, save_optimizer=True): def load_adapter_state(self, model_id, path=None, load_optimizer=True, lr=None): result = self._checkpoint_mgr.load_adapter_state(model_id, path, load_optimizer, lr=lr) self._sync_from_checkpoint_state() + self._sync_registered_lora_session_spec(model_id) return result def save_state(self, checkpoint_path, save_optimizer=True, model_id=None): @@ -1723,6 +2972,9 @@ def sleep(self) -> Dict[str, Any]: Returns: Dict with operation timing information """ + if self._adapter_manager is not None: + raise RuntimeError("sleep is not supported with multi-adapter LoRA server training") + start_time = time.time() # Offload model to CPU @@ -1765,6 +3017,9 @@ def wake_up(self) -> Dict[str, Any]: Returns: Dict with operation timing information """ + if self._adapter_manager is not None: + raise RuntimeError("wake_up is not supported with multi-adapter LoRA server training") + start_time = time.time() device_id = get_device_id() diff --git a/src/xorl/server/runner/runner_dispatcher.py b/src/xorl/server/runner/runner_dispatcher.py index de92fad4..c183007d 100644 --- a/src/xorl/server/runner/runner_dispatcher.py +++ b/src/xorl/server/runner/runner_dispatcher.py @@ -71,6 +71,7 @@ from xorl.data.collators import TextSequenceShardCollator from xorl.distributed.parallel_state import get_parallel_state from xorl.server.protocol.operations import ( + AdapterStateData, EmptyData, LoadStateData, ModelPassData, @@ -191,9 +192,9 @@ def __init__( ) # Operations that participate in cross-rank error sync. - # Only compute ops where all ranks execute in lockstep. - # Excludes weight sync, save/load, etc. that have their own sync mechanisms. - _ERROR_SYNC_OPS = {"forward_backward", "forward", "optim_step"} + # These ops execute on every rank, and rank-0 success is not trustworthy if any + # worker rank failed locally. + _ERROR_SYNC_OPS = {"forward_backward", "forward", "optim_step", "register_session"} def _sync_error_state(self) -> Optional[str]: """Synchronize error state across all ranks via Gloo group. @@ -355,7 +356,7 @@ async def _worker_event_loop(self): # Set error state so rank 0 can detect it during sync self._worker_error = error_msg - # Post-execution error sync for compute ops only (matches rank 0) + # Post-execution error sync for lockstep ops (matches rank 0) if command_type in self._ERROR_SYNC_OPS: try: cross_rank_error = self._sync_error_state() @@ -400,6 +401,7 @@ async def _worker_event_loop(self): "wake_up": "_handle_wake_up", "health_check": "_handle_health_check", "sync_inference_weights": "_handle_sync_inference_weights", + "register_session": "_handle_register_session", "register_adapter": "_handle_register_adapter", "save_adapter_state": "_handle_save_adapter_state", "load_adapter_state": "_handle_load_adapter_state", @@ -465,8 +467,7 @@ async def _handle_request_rank0(self, request: RunnerDispatchCommand) -> RunnerR ) result = await getattr(self, handler_name)(command_dict) - # Post-execution error sync for compute ops only - # (weight sync, save/load have their own distributed sync) + # Post-execution error sync for lockstep ops. if command_type in self._ERROR_SYNC_OPS: cross_rank_error = self._sync_error_state() if cross_rank_error: @@ -490,7 +491,7 @@ async def _handle_request_rank0(self, request: RunnerDispatchCommand) -> RunnerR else: logger.error(f"Rank {self.rank}: Error handling request: {e}", exc_info=True) - # Set error state so other ranks can detect it during compute ops + # Set error state so other ranks can detect it during lockstep ops if command_type in self._ERROR_SYNC_OPS: self._worker_error = str(e) try: @@ -569,12 +570,13 @@ async def _handle_compute_rank0_scatter(self, command_dict: Dict[str, Any], with loss_fn_params = p.loss_fn_params routed_experts = p.routed_experts routed_expert_logits = p.routed_expert_logits - model_id = p.model_id if with_backward else None + model_id = p.model_id or "default" - # Auto-load adapter if it was evicted (all ranks must call this together) + # Auto-load adapter if it was evicted (all ranks must call this together). + # Forward-only requests must also honor model_id so create_model+forward can + # initialize the correct adapter state before save/load operations. was_auto_loaded, auto_load_path = False, None - if with_backward: - was_auto_loaded, auto_load_path = self._adapter_coordinator.auto_load_if_evicted(model_id) + was_auto_loaded, auto_load_path = self._adapter_coordinator.auto_load_if_evicted(model_id) # Select and prepare batches (broadcast-and-select, no scatter) my_batches, routed_experts, routed_expert_logits = self._select_and_prepare_batches( @@ -623,11 +625,10 @@ async def _handle_compute_worker_receive(self, command_dict: Dict[str, Any], wit loss_fn_params = p.loss_fn_params routed_experts = p.routed_experts routed_expert_logits = p.routed_expert_logits - model_id = p.model_id if with_backward else None + model_id = p.model_id or "default" - # Auto-load adapter if it was evicted (all ranks must call this together) - if with_backward: - self._adapter_coordinator.auto_load_if_evicted(model_id) + # Auto-load adapter if it was evicted (all ranks must call this together). + self._adapter_coordinator.auto_load_if_evicted(model_id) # Select and prepare this rank's batches from broadcast data (no scatter) my_batches, routed_experts, routed_expert_logits = self._select_and_prepare_batches( @@ -723,6 +724,38 @@ def _dp_batch_range(dp_rank: int, base_count: int, remainder: int): return dp_rank * (base_count + 1), base_count + 1 return remainder * (base_count + 1) + (dp_rank - remainder) * base_count, base_count + def _batch_parallel_rank_and_size(self, parallel_state, cp_size: int, pp_size: int) -> tuple[int, int]: + """Return the logical data slice rank/size for request batch dispatch. + + Expert-parallel ranks cooperate on one logical batch slice. The + ep_fsdp coordinate identifies which EP group receives a distinct slice; + ranks within that group must all see the same packed batch so MoE/FSDP + collectives and OPD full-vocab KL work stay aligned. + """ + if getattr(parallel_state, "ep_enabled", False): + ep_size = max(1, int(getattr(parallel_state, "ep_size", 1))) + ep_fsdp_size = max(1, int(getattr(parallel_state, "dp_shard_in_ep_size", 1))) + ep_mesh = getattr(parallel_state, "ep_fsdp_device_mesh", None) + if ep_mesh is not None: + try: + ep_fsdp_rank = int(ep_mesh.get_local_rank("ep_fsdp")) + return min(ep_fsdp_rank, ep_fsdp_size - 1), ep_fsdp_size + except Exception as exc: + logger.debug( + "Rank %s: could not read ep_fsdp local rank from EP mesh (%s); falling back to rank arithmetic", + self.rank, + exc, + ) + + ranks_per_pp_stage = max(1, self.world_size // max(1, pp_size)) + local_stage_rank = self.rank % ranks_per_pp_stage + ep_fsdp_rank = min(local_stage_rank // ep_size, ep_fsdp_size - 1) + return ep_fsdp_rank, ep_fsdp_size + + dp_size = max(1, self.world_size // max(1, cp_size * pp_size)) + dp_rank = min(self.rank // max(1, cp_size * pp_size), dp_size - 1) + return dp_rank, dp_size + def _select_and_prepare_batches(self, raw_batches, routed_experts=None, routed_expert_logits=None): """Each rank locally selects its own batches from the full broadcast data. @@ -747,8 +780,7 @@ def _select_and_prepare_batches(self, raw_batches, routed_experts=None, routed_e converted = [self._convert_batch_to_tensors(b) for b in raw_batches] return converted, routed_experts, routed_expert_logits - dp_size = self.world_size // (cp_size * pp_size) - dp_rank = self.rank // (cp_size * pp_size) + dp_rank, dp_size = self._batch_parallel_rank_and_size(parallel_state, cp_size, pp_size) batches_per_dp_group = (num_batches + dp_size - 1) // dp_size if pp_size > 1: batches_per_dp_group = max(batches_per_dp_group, pp_size) @@ -890,6 +922,7 @@ def _execute_compute( my_batches, loss_fn, loss_fn_params, + model_id=model_id, routed_experts=routed_experts, routed_expert_logits=routed_expert_logits, ) @@ -1012,6 +1045,9 @@ async def _handle_save_state(self, command_dict: Dict[str, Any]) -> Dict[str, An f"Rank {self.rank}: Saving state to {checkpoint_path}, model_id={model_id}, save_optimizer={save_optimizer}" ) + if self.trainer.adapter_manager is not None: + self._adapter_coordinator.auto_load_if_evicted(model_id, allow_fresh_materialization=False) + # For LoRA models with save_optimizer=False (sampler weights), use fast PEFT-compatible save # This creates adapter_model.safetensors + adapter_config.json that SGLang can load is_lora_enabled = self.trainer.lora_config.get("enable_lora", False) @@ -1068,6 +1104,9 @@ async def _handle_save_lora_only(self, command_dict: Dict[str, Any]) -> Dict[str logger.debug(f"Rank {self.rank}: Saving LoRA adapter to {lora_path} for model_id={model_id}") + if self.trainer.adapter_manager is not None: + self._adapter_coordinator.auto_load_if_evicted(model_id, allow_fresh_materialization=False) + # NOTE: Cannot use thread pool because trainer.save_lora_only() calls dist.barrier() # which requires all ranks to call from the same thread (main thread). result = self.trainer.save_lora_only(lora_path, model_id=model_id) @@ -1129,17 +1168,28 @@ async def _handle_load_state(self, command_dict: Dict[str, Any]) -> Dict[str, An logger.debug(f"Rank {self.rank}: Checkpoint path exists: {checkpoint_path}") try: - # NOTE: Cannot use thread pool because trainer.load_state() calls dist.barrier() - # which requires all ranks to call from the same thread (main thread). - # This will block the event loop but that's unavoidable for collective operations. - logger.debug(f"Rank {self.rank}: About to call trainer.load_state()...") + # NOTE: Cannot use thread pool because the restore paths use collectives + # that require all ranks to call from the main thread. + if self.trainer.adapter_manager is not None: + logger.debug(f"Rank {self.rank}: Routing load_state through adapter coordinator") + result = await self._adapter_coordinator.handle_load_adapter_state( + { + "payload": AdapterStateData( + model_id=model_id, + path=checkpoint_path, + load_optimizer=load_optimizer, + ) + } + ) + else: + logger.debug(f"Rank {self.rank}: About to call trainer.load_state()...") - # Flush logs before the potentially crashing call - sys.stdout.flush() - sys.stderr.flush() + # Flush logs before the potentially crashing call + sys.stdout.flush() + sys.stderr.flush() - result = self.trainer.load_state(checkpoint_path, load_optimizer, model_id=model_id) - logger.debug(f"Rank {self.rank}: trainer.load_state() returned successfully") + result = self.trainer.load_state(checkpoint_path, load_optimizer, model_id=model_id) + logger.debug(f"Rank {self.rank}: trainer.load_state() returned successfully") # Reset step to 0 after loading state self.trainer.step = 0 @@ -1267,6 +1317,9 @@ def _broadcast_adapter_state(self, model_id: str, default_lr: float) -> None: def _auto_load_adapter_if_evicted(self, model_id: str) -> tuple[bool, str | None]: return self._adapter_coordinator.auto_load_if_evicted(model_id) + async def _handle_register_session(self, command_dict: Dict[str, Any]) -> Dict[str, Any]: + return await self._adapter_coordinator.handle_register_session(command_dict) + async def _handle_register_adapter(self, command_dict: Dict[str, Any]) -> Dict[str, Any]: return await self._adapter_coordinator.handle_register_adapter(command_dict) diff --git a/src/xorl/server/runner/setup.py b/src/xorl/server/runner/setup.py index 8a36d191..29192281 100644 --- a/src/xorl/server/runner/setup.py +++ b/src/xorl/server/runner/setup.py @@ -277,8 +277,10 @@ def main(): break 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://...]") + sys.stderr.write("Error: Config file path required as first positional argument\n") + sys.stderr.write( + "Usage: python -m xorl.server.runner.runner_dispatcher config.yaml [--worker_bind_address tcp://...]\n" + ) sys.exit(1) # Detect config format and parse accordingly diff --git a/src/xorl/server/runner/utils/batch_utils.py b/src/xorl/server/runner/utils/batch_utils.py index bd258182..cb66c1b0 100644 --- a/src/xorl/server/runner/utils/batch_utils.py +++ b/src/xorl/server/runner/utils/batch_utils.py @@ -17,6 +17,31 @@ logger = logging.getLogger(__name__) +def _pad_teacher_hidden_states(value: list[Any]) -> torch.Tensor: + max_len = max(len(seq) for seq in value) + hidden_dim = None + + normalized = [] + for seq in value: + if hasattr(seq, "tolist"): + seq = seq.tolist() + normalized_seq = [] + for hidden in seq: + if hasattr(hidden, "tolist"): + hidden = hidden.tolist() + normalized_seq.append(hidden) + if hidden_dim is None and isinstance(hidden, list): + hidden_dim = len(hidden) + normalized.append(normalized_seq) + + if hidden_dim is None: + raise ValueError("teacher_hidden_states must be a nested sequence of hidden vectors") + + pad_vector = [0.0] * hidden_dim + padded = [seq + [pad_vector] * (max_len - len(seq)) for seq in normalized] + return torch.tensor(padded, dtype=torch.float) + + def convert_batch_to_tensors(batch: Dict[str, Any], rank: int = 0) -> Dict[str, Any]: """ Convert batch data from lists to torch tensors, with padding if needed. @@ -31,7 +56,16 @@ def convert_batch_to_tensors(batch: Dict[str, Any], rank: int = 0) -> Dict[str, converted_batch = {} # Fields that should be float tensors (probabilities, advantages, etc.) - float_fields = {"logprobs", "advantages", "old_logprobs", "values", "returns"} + float_fields = { + "logprobs", + "advantages", + "old_logprobs", + "values", + "returns", + "teacher_weights", + "teacher_hidden_states", + } + long_fields = {"teacher_ids", "teacher_cache_indices"} # Fields that must be int32 (flash attention requires cu_seqlens as int32) int32_fields = {"cu_seq_lens_q", "cu_seq_lens_k"} @@ -43,6 +77,8 @@ def convert_batch_to_tensors(batch: Dict[str, Any], rank: int = 0) -> Dict[str, dtype = torch.float elif key in int32_fields: dtype = torch.int32 + elif key in long_fields: + dtype = torch.long else: dtype = torch.long @@ -56,18 +92,23 @@ def convert_batch_to_tensors(batch: Dict[str, Any], rank: int = 0) -> Dict[str, # If conversion failed (likely due to ragged sequences), try padding 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) - padded = [] - for seq in value: - padded_seq = seq + [pad_value] * (max_len - len(seq)) - padded.append(padded_seq) try: - # Determine dtype for padded sequences - dtype = torch.float if key in float_fields else torch.long - tensor = torch.tensor(padded, dtype=dtype) + if key == "teacher_hidden_states": + tensor = _pad_teacher_hidden_states(value) + max_len = tensor.shape[1] + dtype = torch.float + else: + 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) + padded = [] + for seq in value: + padded_seq = seq + [pad_value] * (max_len - len(seq)) + padded.append(padded_seq) + # Determine dtype for padded sequences + dtype = torch.float if key in float_fields else torch.long + tensor = torch.tensor(padded, dtype=dtype) converted_batch[key] = tensor logger.debug( f"Rank {rank}: Padded and converted {key}: {len(value)} sequences, max_len={max_len}, dtype={dtype}" @@ -169,23 +210,25 @@ def simple_sequence_shard(batch: Dict[str, Any], rank: int = 0) -> Dict[str, Any pad_len = cp_chunk_size * cp_size - seq_len # Helper to pad and slice tensors - def pad_and_slice(tensor, pad_value=0): + def pad_and_slice(tensor, pad_value=0, seq_dim=-1): if tensor is None: return None if not isinstance(tensor, torch.Tensor): tensor = torch.tensor(tensor, dtype=torch.long) + if seq_dim < 0: + seq_dim = tensor.ndim + seq_dim + # Pad if needed if pad_len > 0: pad_shape = list(tensor.shape) - pad_shape[-1] = pad_len + pad_shape[seq_dim] = pad_len pad_tensor = torch.full(pad_shape, pad_value, dtype=tensor.dtype, device=tensor.device) - tensor = torch.cat([tensor, pad_tensor], dim=-1) + tensor = torch.cat([tensor, pad_tensor], dim=seq_dim) # Slice for this cp_rank start_idx = cp_rank * cp_chunk_size - end_idx = start_idx + cp_chunk_size - return tensor[..., start_idx:end_idx] + return tensor.narrow(seq_dim, start_idx, cp_chunk_size) # Apply to all sequence tensors sharded_batch = {} @@ -212,6 +255,14 @@ def pad_and_slice(tensor, pad_value=0): sharded_batch[key] = torch.cat([value, pad_tensor], dim=-1) else: sharded_batch[key] = value + elif key == "teacher_hidden_states": + if not isinstance(value, torch.Tensor): + value = torch.tensor(value, dtype=torch.float) + elif not torch.is_floating_point(value): + value = value.float() + if value.dim() == 2: + value = value.unsqueeze(0) + sharded_batch[key] = pad_and_slice(value, pad_value=0.0, seq_dim=1) elif isinstance(value, torch.Tensor) and value.dim() >= 1 and value.size(-1) == seq_len: # Other tensors with matching sequence length # Use appropriate pad value based on field type diff --git a/src/xorl/server/runner/utils/routing_replay_handler.py b/src/xorl/server/runner/utils/routing_replay_handler.py index 407d95fb..c107d45e 100644 --- a/src/xorl/server/runner/utils/routing_replay_handler.py +++ b/src/xorl/server/runner/utils/routing_replay_handler.py @@ -23,7 +23,6 @@ import base64 import logging -import math from typing import Any, Dict, List, Optional, Union import numpy as np @@ -68,10 +67,29 @@ def __init__(self, model: nn.Module) -> None: @staticmethod def _extract_topk(model: nn.Module) -> Optional[int]: - """Extract num_experts_per_tok from the model config, if available.""" + """Extract num_experts_per_tok from the model config, if available. + + Qwen3.6 nests `num_experts_per_tok` under `config.text_config`, so we + check the nested config first to avoid silently picking up the wrong + top-k width from row 0 of mixed routing data. + """ config = getattr(model, "config", None) + configs = [config] if config is not None: - topk = getattr(config, "num_experts_per_tok", None) + text_config = ( + config.get("text_config") if isinstance(config, dict) else getattr(config, "text_config", None) + ) + if text_config is not None: + configs.insert(0, text_config) + + for candidate in configs: + if candidate is None: + continue + topk = ( + candidate.get("num_experts_per_tok") + if isinstance(candidate, dict) + else getattr(candidate, "num_experts_per_tok", None) + ) if topk is not None: return int(topk) return None @@ -255,9 +273,11 @@ def fill_routing_replay( logger.warning("R3: No valid routing data after decoding") return False - # Infer dimensions from first decoded datum + # Prefer the model config top-k when available, then fall back to the + # first decoded datum. Mixed 6/8 routing rows (Qwen3.6) can otherwise + # let row 0 pick the wrong width and crash tensorization downstream. num_layers_in_data = len(decoded_routing[0][0]) - topk = len(decoded_routing[0][0][0]) + topk = self._model_topk or len(decoded_routing[0][0][0]) total_tokens_raw = sum(len(d) for d in decoded_routing) logger.debug( @@ -323,11 +343,115 @@ def _build_per_mb_routing( Returns: List of tensors, each [num_tokens_mb, num_layers, topk] as torch.long. """ - cp_enabled = get_parallel_state().cp_enabled + parallel_state = get_parallel_state() + cp_enabled = parallel_state.cp_enabled if cp_enabled: - cp_size = get_parallel_state().cp_size - cp_rank = get_parallel_state().cp_rank - pad_to_multiple_of = math.lcm(128, cp_size) + cp_size = parallel_state.cp_size + cp_rank = parallel_state.cp_rank + + def _pad_entry(): + return [list(range(topk)) for _ in range(num_layers_in_data)] + + def _num_tokens(value: Any) -> Optional[int]: + if value is None: + return None + if isinstance(value, torch.Tensor): + return int(value.numel()) + if isinstance(value, list): + if value and isinstance(value[0], list): + return sum(len(row) if isinstance(row, list) else 1 for row in value) + return len(value) + return None + + def _first_dim(value: Any) -> Optional[int]: + if isinstance(value, torch.Tensor) and value.ndim >= 2: + return int(value.shape[0]) + if isinstance(value, list) and value and isinstance(value[0], list): + return len(value) + return None + + def _last_dim(value: Any) -> Optional[int]: + if isinstance(value, torch.Tensor) and value.ndim >= 1: + return int(value.shape[-1]) + if isinstance(value, list): + if value and isinstance(value[0], list): + return len(value[0]) + return len(value) + return None + + def _flatten_position_ids(value: Any) -> Optional[List[int]]: + if value is None: + return None + if isinstance(value, torch.Tensor): + return [int(v) for v in value.reshape(-1).tolist()] + if isinstance(value, list): + if value and isinstance(value[0], list): + return [int(v) for row in value for v in row] + return [int(v) for v in value] + return None + + def _resize_position_ids(position_ids: List[int], target_tokens: int) -> List[int]: + if len(position_ids) < target_tokens: + pad_count = target_tokens - len(position_ids) + position_ids = position_ids + [i % 1024 for i in range(pad_count)] + elif len(position_ids) > target_tokens: + position_ids = position_ids[:target_tokens] + return position_ids + + def _zigzag_reorder_routing( + routing: List[Any], + position_ids: List[int], + ringattn_size: int, + mb_idx: int, + ) -> List[Any]: + if ringattn_size <= 1: + return routing + + boundaries = [i for i, pos in enumerate(position_ids) if pos == 0] + boundaries.append(len(position_ids)) + num_subchunks = 2 * ringattn_size + rank_parts = [[] for _ in range(ringattn_size)] + + for boundary_idx in range(len(boundaries) - 1): + start_idx = boundaries[boundary_idx] + end_idx = boundaries[boundary_idx + 1] + doc_len = end_idx - start_idx + if doc_len == 0: + continue + if doc_len % num_subchunks != 0: + raise ValueError( + f"R3: MB{mb_idx} document at position {start_idx} has length {doc_len}, " + f"not divisible by 2*ringattn_size={num_subchunks}. " + "Routing replay cannot match ring-attention zigzag layout." + ) + + subchunk_len = doc_len // num_subchunks + chunks = [ + routing[start_idx + subchunk_idx * subchunk_len : start_idx + (subchunk_idx + 1) * subchunk_len] + for subchunk_idx in range(num_subchunks) + ] + for ring_rank in range(ringattn_size): + rank_parts[ring_rank].extend(chunks[ring_rank]) + rank_parts[ring_rank].extend(chunks[num_subchunks - 1 - ring_rank]) + + return [token_routing for rank_part in rank_parts for token_routing in rank_part] + + def _resize_routing(routing: List[Any], target_tokens: Optional[int], mb_idx: int, reason: str) -> List[Any]: + if target_tokens is None: + return routing + if len(routing) < target_tokens: + pad_count = target_tokens - len(routing) + logger.debug("R3: Padded MB%s routing by %s tokens to match %s", mb_idx, pad_count, reason) + return routing + [_pad_entry()] * pad_count + if len(routing) > target_tokens: + logger.debug( + "R3: Truncated MB%s routing by %s tokens to match %s", + mb_idx, + len(routing) - target_tokens, + reason, + ) + return routing[:target_tokens] + return routing # Read num_samples from each micro-batch (set by SequentialPacker._finalize_packed_batch) micro_batch_datum_counts = [mb.get("num_samples", 1) for mb in micro_batches] @@ -350,64 +474,93 @@ def _build_per_mb_routing( for mb_idx, num_datums in enumerate(micro_batch_datum_counts): # Concatenate routing from all datums packed into this micro-batch - mb_routing = [] + datum_routing = [] for _ in range(num_datums): if datum_cursor < len(decoded_routing): - mb_routing.extend(decoded_routing[datum_cursor]) + datum_routing.append(decoded_routing[datum_cursor]) datum_cursor += 1 + mb_routing = [token_routing for datum in datum_routing for token_routing in datum] mb_total_tokens = len(mb_routing) + micro_batch = micro_batches[mb_idx] if mb_idx < len(micro_batches) else {} + expected_mb_tokens = _num_tokens(micro_batch.get("input_ids")) if cp_enabled and mb_total_tokens > 0: - # CRITICAL: Pad to match packer's pad_to_multiple_of BEFORE SP chunking. - # The SequentialPacker pads sequences to lcm(128, cp_size) before - # TextSequenceShardCollator shards them. We must replicate this step - # so routing SP-slice boundaries match the actual data boundaries. - packing_pad = (pad_to_multiple_of - mb_total_tokens % pad_to_multiple_of) % pad_to_multiple_of - if packing_pad > 0: - pad_entry = [list(range(topk)) for _ in range(num_layers_in_data)] - mb_routing = mb_routing + [pad_entry] * packing_pad - - padded_len = mb_total_tokens + packing_pad - - # SP-slice the padded block (matches TextSequenceShardCollator.sp_slice) - cp_chunk_size = (padded_len + cp_size - 1) // cp_size - sp_pad_count = cp_chunk_size * cp_size - padded_len - - if sp_pad_count > 0: - pad_entry = [list(range(topk)) for _ in range(num_layers_in_data)] - mb_routing = mb_routing + [pad_entry] * sp_pad_count + # Match the actual sharded micro-batch shape. Packed batches may + # already be padded to 128-token boundaries by SequentialPacker, + # while unpacked/server batches are only padded to the CP size. + # position_ids stays full-length after sequence sharding, so it is + # the source of truth for how much routing data existed before the + # local CP slice. + input_ids = micro_batch.get("input_ids") + position_ids = micro_batch.get("position_ids") + batch_rows = _first_dim(input_ids) + local_seq_len = _last_dim(input_ids) + full_seq_len = _last_dim(position_ids) + ringattn_size = getattr(parallel_state, "ringattn_size", 1) + rowwise_unpacked = ( + batch_rows is not None + and batch_rows > 1 + and batch_rows == len(datum_routing) + and local_seq_len is not None + and full_seq_len is not None + ) - start = cp_rank * cp_chunk_size - end = (cp_rank + 1) * cp_chunk_size - mb_routing = mb_routing[start:end] + if rowwise_unpacked: + sharded_routing = [] + cp_chunk_size = local_seq_len + start = cp_rank * cp_chunk_size + end = start + cp_chunk_size + for row_routing in datum_routing: + row_routing = _resize_routing(list(row_routing), full_seq_len, mb_idx, "full row length") + sharded_routing.extend(row_routing[start:end]) + mb_routing = sharded_routing + logger.debug( + "R3: SP MB%s rowwise - raw_tokens=%s, rows=%s, full_seq=%s, local_seq=%s, cp_rank=%s", + mb_idx, + mb_total_tokens, + batch_rows, + full_seq_len, + local_seq_len, + cp_rank, + ) + else: + full_tokens = _num_tokens(position_ids) + if full_tokens is None: + full_tokens = ((mb_total_tokens + cp_size - 1) // cp_size) * cp_size + mb_routing = _resize_routing(mb_routing, full_tokens, mb_idx, "full SP-padded length") + if ringattn_size > 1: + zigzag_position_ids = _flatten_position_ids(micro_batch.get("_original_position_ids")) + if zigzag_position_ids is None: + zigzag_position_ids = _flatten_position_ids(position_ids) + if zigzag_position_ids is None: + logger.warning( + "R3: MB%s ring-attention routing lacks position_ids; falling back to contiguous slice", + mb_idx, + ) + else: + zigzag_position_ids = _resize_position_ids(zigzag_position_ids, len(mb_routing)) + mb_routing = _zigzag_reorder_routing( + mb_routing, + zigzag_position_ids, + ringattn_size, + mb_idx, + ) + cp_chunk_size = expected_mb_tokens or ((len(mb_routing) + cp_size - 1) // cp_size) + start = cp_rank * cp_chunk_size + end = start + cp_chunk_size + mb_routing = mb_routing[start:end] logger.debug( f"R3: SP MB{mb_idx} - {mb_total_tokens} tokens ({num_datums} datums), " - f"packing_pad={packing_pad}, padded_len={padded_len}, " - f"sp_chunk={cp_chunk_size}, sp_pad={sp_pad_count}, slice [{start}:{end}]" + f"sp_chunk={cp_chunk_size}, ringattn_size={ringattn_size}, slice [{start}:{end}]" ) # Pad routing to match actual micro-batch token count. # The packer's pad_to_multiple_of may have added padding tokens # that aren't in the raw routing data. - if mb_idx < len(micro_batches) and "input_ids" in micro_batches[mb_idx]: - mb_input_ids = micro_batches[mb_idx]["input_ids"] - if isinstance(mb_input_ids, torch.Tensor): - expected_mb_tokens = mb_input_ids.shape[0] * mb_input_ids.shape[1] - else: - expected_mb_tokens = ( - len(mb_input_ids[0]) if isinstance(mb_input_ids[0], list) else len(mb_input_ids) - ) - - if len(mb_routing) < expected_mb_tokens: - pad_count = expected_mb_tokens - len(mb_routing) - pad_entry = [list(range(topk)) for _ in range(num_layers_in_data)] - mb_routing.extend([pad_entry] * pad_count) - logger.debug( - f"R3: Padded MB{mb_idx} routing by {pad_count} tokens to match " - f"micro-batch size ({expected_mb_tokens})" - ) + if expected_mb_tokens is not None: + mb_routing = _resize_routing(mb_routing, expected_mb_tokens, mb_idx, "micro-batch size") if mb_routing: # Convert to tensor: [num_tokens_mb, num_layers, topk] diff --git a/src/xorl/server/server_arguments.py b/src/xorl/server/server_arguments.py index c2f452ab..6f76d583 100644 --- a/src/xorl/server/server_arguments.py +++ b/src/xorl/server/server_arguments.py @@ -11,8 +11,11 @@ from dataclasses import dataclass, field from typing import Any, Dict, List, Literal, Optional +import torch import yaml +from xorl.ops.loss import CrossEntropyMode + @dataclass class ServerArguments: @@ -89,6 +92,16 @@ class ServerArguments: }, ) + record_routing_weights: bool = field( + default=True, + metadata={ + "help": "Cache routing weights on the forward pass so they can override the " + "regathered weights during checkpoint recompute. Needed only when the " + "attention forward is non-deterministic across recompute. Disabling skips the " + "per-layer pinned CPU allocation + D2H/H2D copies on every step." + }, + ) + deepep_buffer_size_gb: float = field( default=2.0, metadata={"help": "DeepEP buffer size in GB (effective when ep_dispatch='deepep')."} ) @@ -131,6 +144,11 @@ class ServerArguments: default=False, metadata={"help": "Explicitly cast Q/K to bfloat16 after RoPE for SGLang alignment."} ) + flash_attention_deterministic: bool = field( + default=False, + metadata={"help": "Request FlashAttention deterministic backward kernels when available."}, + ) + # Multimodal model configuration foundation: Dict[str, str] = field(default_factory=dict, metadata={"help": "Foundation model extra config"}) @@ -194,6 +212,8 @@ class ServerArguments: seed: int = field(default=42, metadata={"help": "Random seed for reproducibility"}) + enable_full_determinism: bool = field(default=False, metadata={"help": "Enable full deterministic execution."}) + 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"}) @@ -202,6 +222,16 @@ class ServerArguments: enable_activation_offload: bool = field(default=False, metadata={"help": "Enable activation CPU offloading"}) + activation_gpu_limit: float = field( + default=0.0, + metadata={ + "help": ( + "When enabling activation offload, the number of GB of activations allowed to remain on GPU. " + "Defaults to 0.0, which offloads all eligible activations." + ) + }, + ) + enable_compile: bool = field(default=False, metadata={"help": "Enable torch.compile for model forward pass"}) enable_reentrant: bool = field( @@ -217,14 +247,15 @@ class ServerArguments: ) load_weights_mode: str = field( - default="auto", metadata={"help": "Weight loading mode: 'auto', 'safetensors', 'dcp'"} + default="grouped", + metadata={"help": ("Weight loading mode: 'grouped' (default, with rank-0 fallback), 'all_ranks', or 'skip'")}, ) init_device: Optional[Literal["cpu", "meta", "cuda"]] = field( default="meta", metadata={"help": "Device for model initialization"} ) - ce_mode: Literal["eager", "compiled"] = field( + ce_mode: CrossEntropyMode = field( default="compiled", metadata={ "help": "Cross-entropy implementation: 'compiled' (RECOMMENDED, torch.compile) or 'eager' (baseline, may OOM at 32K)" @@ -235,19 +266,39 @@ class ServerArguments: # Optimizer # ======================================================================== - optimizer: Literal["adamw", "anyprecision_adamw", "sgd", "signsgd", "muon"] = field( + optimizer: Literal["adamw", "anyprecision_adamw", "sgd", "signsgd", "distsignsgd", "muon"] = field( default="adamw", metadata={ - "help": "Optimizer type. 'signsgd' is a state-free sign update; 'muon' uses " + "help": "Optimizer type. 'signsgd' is a local state-free sign update; " + "'distsignsgd' signs gradients before FSDP2 reduction; 'muon' uses " "Newton-Schulz orthogonalization for 2D+ weight matrices." }, ) + lr: float = field( + default=1e-5, + metadata={"help": "Default learning rate for the server's implicit/default training session."}, + ) + + weight_decay: float = field( + default=0.01, + metadata={"help": "Default weight decay for the server's implicit/default training session."}, + ) + optimizer_dtype: Literal["fp32", "bf16"] = field( default="bf16", metadata={"help": "Dtype for optimizer states (momentum/variance). 'bf16' halves optimizer memory."}, ) + cautious_weight_decay: bool = field( + default=False, + metadata={ + "help": "Apply Cautious Weight Decay (Chen et al., arXiv:2510.12402): " + "mask the decoupled decay term by I(u_t * x_t >= 0). With optimizer='adamw' " + "this routes to AnyPrecisionAdamW with fp32 state (no fused kernel)." + }, + ) + muon_lr: float = field( default=0.02, metadata={ @@ -279,10 +330,11 @@ class ServerArguments: ) muon_ns_algorithm: Literal["standard_newton_schulz", "gram_newton_schulz"] = field( - default="standard_newton_schulz", + default="gram_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." + "help": "Newton-Schulz backend for Muon. 'gram_newton_schulz' (default) batches across " + "MoE experts via baddbmm and is ~2x faster on Qwen3.5-style MoE; 'standard_newton_schulz' " + "uses the PyTorch upstream path for bit-exact equivalence with torch.optim._muon." }, ) @@ -309,6 +361,13 @@ class ServerArguments: "A value of 2 means restart after the second iteration." }, ) + muon_grouped_gram_ns_fp32_byte_limit: int = field( + default=512 * 1024**2, + metadata={ + "help": "Maximum fp32 scratch bytes per grouped Muon Gram Newton-Schulz batch before chunking. " + "Lower values reduce peak optimizer scratch memory at the cost of more launches." + }, + ) muon_grad_dtype: Optional[Literal["fp32", "bf16"]] = field( default=None, metadata={ @@ -331,6 +390,28 @@ class ServerArguments: }, ) + muon_distributed_mode: Literal["shard_local", "full_gradient"] = field( + default="shard_local", + metadata={ + "help": "How Muon handles Newton-Schulz on FSDP2/EP-sharded DTensor params. " + "'shard_local': run NS on each rank's local shard (cheap, approximate). " + "'full_gradient': all-gather post-momentum update, run NS on the full matrix on " + "every rank in the param's mesh, slice back to the local shard. Implements the " + "dense path of DeepSeek V4 §3.5.1." + }, + ) + + moe_grad_reduce_mode: Literal["reduce_scatter", "bf16_a2a_fp32_sum"] = field( + default="reduce_scatter", + metadata={ + "help": "Reduce-scatter strategy for MoE expert gradients on the ep_fsdp mesh dim. " + "'reduce_scatter': default NCCL reduce-scatter. " + "'bf16_a2a_fp32_sum': stochastic-round FP32 grads to BF16, all-to-all across the " + "ep_fsdp group, then sum the per-rank chunks locally in FP32. Halves comm volume " + "while preserving FP32 accumulation. Implements the MoE path of DeepSeek V4 §3.5.1." + }, + ) + # ======================================================================== # Checkpointing & Output # ======================================================================== @@ -432,6 +513,14 @@ class ServerArguments: lora_rank: int = field(default=32, metadata={"help": "LoRA rank (r parameter)"}) + max_lora_rank: Optional[int] = field( + default=None, + metadata={ + "help": "Maximum LoRA rank allocated in the server model substrate. Defaults to lora_rank. " + "Per-session ranks must be <= max_lora_rank." + }, + ) + lora_alpha: int = field(default=16, metadata={"help": "LoRA alpha scaling parameter"}) lora_target_modules: Optional[List[str]] = field( @@ -448,6 +537,13 @@ class ServerArguments: }, ) + lora_export_format: str = field( + default="peft", + metadata={ + "help": "On-disk layout for MoE LoRA export. 'peft' (default) writes per-expert keys in PEFT orientation. 'sglang_shared_outer' writes stacked 3D tensors under experts.w{1,2,3} in SGLang's shared_outer format (requires moe_hybrid_shared_lora=True)." + }, + ) + # ======================================================================== # QLoRA Configuration # ======================================================================== @@ -470,6 +566,13 @@ class ServerArguments: reset_optimizer_on_merge: bool = field( default=False, metadata={"help": "ReLoRA-style partial optimizer reset after each LoRA merge"} ) + adapter_state_load_mode: Literal["all_ranks", "rank0_broadcast"] = field( + default="all_ranks", + metadata={ + "help": "How to restore multi-adapter LoRA checkpoints. 'all_ranks': each rank loads adapter state locally. " + "'rank0_broadcast': rank 0 loads once and broadcasts weights, metadata, and optimizer state." + }, + ) # ======================================================================== # MoE Training Configuration @@ -481,14 +584,49 @@ class ServerArguments: # Inference Weight Sync Configuration # ======================================================================== - sync_inference_method: Literal["nccl_broadcast"] = field( + sync_inference_method: Literal["nccl_broadcast", "nccl_simple", "p2p"] = field( default="nccl_broadcast", metadata={ "help": "Method for syncing weights to inference endpoints: " - "'nccl_broadcast' (rank-0 broadcast via SGLang update_weights_from_distributed)" + "'nccl_broadcast' (rank-0 broadcast via SGLang update_weights_from_distributed, " + "interleaved with the FSDP unshard loop); " + "'nccl_simple' (two-phase: stage all params to CPU during the FSDP loop, then " + "broadcast in chunks — FSDP and weight-sync NCCL communicators never interleave); " + "'p2p' (RDMA one-sided writes via Mooncake TransferEngine into SGLang's " + "registered param memory; requires --enable-rdma-weight-updates on the SGLang side)" }, ) + @property + def optimizer_kwargs(self) -> Dict[str, Any]: + """Collect optimizer-specific kwargs for build_optimizer.""" + kwargs: Dict[str, Any] = {} + if self.optimizer == "muon": + kwargs["muon_lr"] = self.muon_lr + kwargs["muon_momentum"] = self.muon_momentum + 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 + kwargs["muon_gram_ns_restart_iterations"] = self.muon_gram_ns_restart_iterations + 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 + if self.muon_distributed_mode != "shard_local": + kwargs["muon_distributed_mode"] = self.muon_distributed_mode + return kwargs + def __post_init__(self): """Validate and set defaults.""" # Set default paths @@ -512,6 +650,37 @@ def __post_init__(self): # the launcher can still use them via engine_connect_host + worker_bind_port pass + if self.adapter_state_load_mode not in {"all_ranks", "rank0_broadcast"}: + raise ValueError( + "adapter_state_load_mode must be 'all_ranks' or 'rank0_broadcast', " + f"got {self.adapter_state_load_mode!r}" + ) + if self.enable_lora and self.pipeline_parallel_size > 1: + raise ValueError( + "pipeline_parallel_size > 1 is not supported with multi-adapter LoRA server training. " + "Adapter coordination currently assumes identical local LoRA layouts on every rank." + ) + if self.enable_lora and self.merge_lora_interval > 0: + raise ValueError("merge_lora_interval is not supported with multi-adapter LoRA server training") + if self.max_lora_rank is None: + self.max_lora_rank = self.lora_rank + if self.max_lora_rank < self.lora_rank: + raise ValueError( + f"max_lora_rank ({self.max_lora_rank}) must be >= lora_rank ({self.lora_rank}) for the default session" + ) + + if self.load_weights_mode not in {"grouped", "all_ranks", "skip"}: + raise ValueError( + f"Unsupported load_weights_mode={self.load_weights_mode!r}. Expected one of: grouped, all_ranks, skip." + ) + + if self.load_weights_mode == "skip" and not self.load_checkpoint_path: + raise ValueError( + "load_weights_mode='skip' skips HF weight loading and relies on " + "load_checkpoint_path to materialize parameters from a DCP checkpoint. " + "Set load_checkpoint_path or choose a different load_weights_mode." + ) + def to_config_dict(self) -> Dict[str, Any]: """ Convert ServerArguments to the config dict format expected by ModelRunner. @@ -529,6 +698,7 @@ def to_config_dict(self) -> Dict[str, Any]: "moe_implementation": self.moe_implementation, "ep_dispatch": self.ep_dispatch, "train_router": self.train_router, + "record_routing_weights": self.record_routing_weights, "deepep_buffer_size_gb": self.deepep_buffer_size_gb, "deepep_num_sms": self.deepep_num_sms, "deepep_async_combine": self.deepep_async_combine, @@ -542,10 +712,12 @@ def to_config_dict(self) -> Dict[str, Any]: "activation_native": self.activation_native, "rope_native": self.rope_native, "attention_cast_bf16": self.attention_cast_bf16, + "flash_attention_deterministic": self.flash_attention_deterministic, }, "train": { "output_dir": self.output_dir, "seed": self.seed, + "enable_full_determinism": self.enable_full_determinism, "data_parallel_mode": self.data_parallel_mode, "ulysses_parallel_size": self.ulysses_parallel_size, "expert_parallel_size": self.expert_parallel_size, @@ -558,6 +730,7 @@ def to_config_dict(self) -> Dict[str, Any]: "enable_gradient_checkpointing": self.enable_gradient_checkpointing, "enable_full_shard": self.enable_full_shard, "enable_activation_offload": self.enable_activation_offload, + "activation_gpu_limit": self.activation_gpu_limit, "enable_compile": self.enable_compile, "enable_reentrant": self.enable_reentrant, "enable_forward_prefetch": self.enable_forward_prefetch, @@ -566,7 +739,10 @@ def to_config_dict(self) -> Dict[str, Any]: "init_device": self.init_device, "ce_mode": self.ce_mode, "optimizer": self.optimizer, + "lr": self.lr, + "weight_decay": self.weight_decay, "optimizer_dtype": self.optimizer_dtype, + "cautious_weight_decay": self.cautious_weight_decay, "muon_lr": self.muon_lr, "muon_momentum": self.muon_momentum, "muon_nesterov": self.muon_nesterov, @@ -576,9 +752,13 @@ def to_config_dict(self) -> Dict[str, Any]: "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_grouped_gram_ns_fp32_byte_limit": self.muon_grouped_gram_ns_fp32_byte_limit, "muon_grad_dtype": self.muon_grad_dtype, "muon_update_dtype": self.muon_update_dtype, "muon_force_momentum_path": self.muon_force_momentum_path, + "muon_distributed_mode": self.muon_distributed_mode, + "moe_grad_reduce_mode": self.moe_grad_reduce_mode, + "optimizer_kwargs": self.optimizer_kwargs, "load_checkpoint_path": self.load_checkpoint_path, "ckpt_manager": self.ckpt_manager, "enable_self_test": self.enable_self_test, @@ -598,15 +778,18 @@ def to_config_dict(self) -> Dict[str, Any]: "lora": { "enable_lora": self.enable_lora, "lora_rank": self.lora_rank, + "max_lora_rank": self.max_lora_rank, "lora_alpha": self.lora_alpha, "lora_target_modules": self.lora_target_modules, "moe_hybrid_shared_lora": self.moe_hybrid_shared_lora, + "lora_export_format": self.lora_export_format, "enable_qlora": self.enable_qlora, "quant_format": self.quant_format, "quant_group_size": self.quant_group_size, "exclude_modules": self.qlora_exclude_modules, "merge_lora_interval": self.merge_lora_interval, "reset_optimizer_on_merge": self.reset_optimizer_on_merge, + "adapter_state_load_mode": self.adapter_state_load_mode, }, } return config diff --git a/src/xorl/server/session_spec.py b/src/xorl/server/session_spec.py new file mode 100644 index 00000000..bc52067d --- /dev/null +++ b/src/xorl/server/session_spec.py @@ -0,0 +1,445 @@ +"""Utilities for multi-adapter session runtime specs. + +The server supports heterogeneous multi-adapter LoRA sessions where each +``model_id`` owns: + +- a LoRA runtime config (rank + alpha) +- an optimizer contract (type + kwargs + default learning rate) + +This module normalizes those specs into a shared JSON-safe structure used by +the API server, worker runtime, and checkpoint metadata. +""" + +from __future__ import annotations + +import json +import os +from copy import deepcopy +from typing import Any, Dict, Optional + +import torch + + +SUPPORTED_OPTIMIZER_TYPES = { + "adamw", + "anyprecision_adamw", + "sgd", + "signsgd", + "muon", +} + +DEFAULT_ADAM_BETAS = (0.9, 0.95) +DEFAULT_ADAM_EPS = 1e-8 +SESSION_SPEC_FILENAME = "session_spec.json" + +_PER_SESSION_LORA_KEY_ALIASES = { + "rank": "lora_rank", + "alpha": "lora_alpha", + "target_modules": "lora_target_modules", +} + +_SERVER_WIDE_LORA_KEYS = { + "enable_lora", + "enable_qlora", + "lora_target_modules", + "train_attn", + "train_mlp", + "train_unembed", + "moe_shared_lora", + "moe_hybrid_shared_lora", + "quant_format", + "quant_group_size", + "exclude_modules", +} + +_DTYPE_STRING_TO_TORCH = { + "bf16": torch.bfloat16, + "fp32": torch.float32, +} + + +def _clone_jsonable(value: Any) -> Any: + """Deep copy nested metadata while making torch dtypes JSON-safe.""" + if isinstance(value, torch.dtype): + if value == torch.bfloat16: + return "bf16" + if value == torch.float32: + return "fp32" + return str(value) + if isinstance(value, dict): + return {k: _clone_jsonable(v) for k, v in value.items()} + if isinstance(value, list): + return [_clone_jsonable(v) for v in value] + if isinstance(value, tuple): + return [_clone_jsonable(v) for v in value] + return deepcopy(value) + + +def _restore_optimizer_metadata(value: Any) -> Any: + """Restore JSON-safe optimizer kwargs to the build_optimizer format.""" + if isinstance(value, dict): + restored = {} + for key, nested in value.items(): + converted = _restore_optimizer_metadata(nested) + if key in {"muon_momentum_dtype", "muon_grad_dtype", "muon_update_dtype"} and isinstance(converted, str): + converted = _DTYPE_STRING_TO_TORCH.get(converted, converted) + restored[key] = converted + return restored + if isinstance(value, list): + return [_restore_optimizer_metadata(v) for v in value] + return value + + +def _normalize_lora_config_keys(raw_lora_config: Optional[Dict[str, Any]]) -> Dict[str, Any]: + data = dict(raw_lora_config or {}) + normalized: Dict[str, Any] = {} + for key, value in data.items(): + normalized[_PER_SESSION_LORA_KEY_ALIASES.get(key, key)] = value + return normalized + + +def normalize_lora_runtime_config( + raw_lora_config: Optional[Dict[str, Any]], + *, + default_rank: int, + default_alpha: int, + max_lora_rank: int, + server_lora_config: Optional[Dict[str, Any]] = None, +) -> Dict[str, Any]: + """Normalize session LoRA config to the supported per-session surface.""" + lora_config = _normalize_lora_config_keys(raw_lora_config) + server_lora_config = dict(server_lora_config or {}) + + for key in sorted(set(lora_config) - {"lora_rank", "lora_alpha"}): + server_value = server_lora_config.get(key) + if lora_config[key] != server_value: + raise ValueError( + "Per-session LoRA config may only override rank and alpha. " + f"Unsupported override for {key!r}: {lora_config[key]!r} (server={server_value!r})." + ) + + lora_rank = int(lora_config.get("lora_rank", default_rank)) + lora_alpha = int(lora_config.get("lora_alpha", default_alpha)) + + if lora_rank <= 0: + raise ValueError(f"lora_rank must be positive, got {lora_rank}") + if lora_alpha <= 0: + raise ValueError(f"lora_alpha must be positive, got {lora_alpha}") + if lora_rank > max_lora_rank: + raise ValueError( + f"Requested lora_rank={lora_rank} exceeds server max_lora_rank={max_lora_rank}. " + "Increase server.max_lora_rank to support this session." + ) + + return { + "lora_rank": lora_rank, + "lora_alpha": lora_alpha, + } + + +def normalize_optimizer_config( + raw_optimizer_config: Optional[Dict[str, Any]], + *, + default_type: str, + default_learning_rate: float, + default_weight_decay: float, + default_optimizer_dtype: str, + default_optimizer_kwargs: Optional[Dict[str, Any]] = None, + default_betas: tuple[float, float] = DEFAULT_ADAM_BETAS, + default_eps: float = DEFAULT_ADAM_EPS, +) -> Dict[str, Any]: + """Normalize session optimizer config to a JSON-safe runtime contract.""" + raw = dict(raw_optimizer_config or {}) + raw_optimizer_kwargs = _clone_jsonable(raw.get("optimizer_kwargs", default_optimizer_kwargs or {})) + if not isinstance(raw_optimizer_kwargs, dict): + raise ValueError(f"optimizer_config.optimizer_kwargs must be a dict, got {type(raw_optimizer_kwargs)!r}") + optimizer_kwargs = dict(raw_optimizer_kwargs) + + optimizer_type = raw.get("type", default_type) + if optimizer_type not in SUPPORTED_OPTIMIZER_TYPES: + raise ValueError( + f"Unsupported optimizer type {optimizer_type!r}. Supported: {sorted(SUPPORTED_OPTIMIZER_TYPES)}" + ) + + kwargs_learning_rate = optimizer_kwargs.pop("learning_rate", None) + kwargs_lr = optimizer_kwargs.pop("lr", None) + learning_rate = float( + raw.get("learning_rate", raw.get("lr", kwargs_learning_rate or kwargs_lr or default_learning_rate)) + ) + + kwargs_weight_decay = optimizer_kwargs.pop("weight_decay", None) + weight_decay = float( + raw.get("weight_decay", kwargs_weight_decay if kwargs_weight_decay is not None else default_weight_decay) + ) + optimizer_dtype = str(raw.get("optimizer_dtype", default_optimizer_dtype)) + + kwargs_betas = optimizer_kwargs.pop("betas", None) + kwargs_adamw_betas = optimizer_kwargs.pop("adamw_betas", None) + betas_value = raw.get("betas", kwargs_adamw_betas if kwargs_adamw_betas is not None else kwargs_betas) + if betas_value is None: + betas_value = default_betas + + if betas_value is not None: + if not isinstance(betas_value, (list, tuple)) or len(betas_value) != 2: + raise ValueError(f"optimizer_config.betas must be a length-2 list/tuple, got {betas_value!r}") + betas = [float(betas_value[0]), float(betas_value[1])] + else: + betas = None + + kwargs_eps = optimizer_kwargs.pop("eps", None) + kwargs_adamw_eps = optimizer_kwargs.pop("adamw_eps", None) + eps_value = raw.get("eps", kwargs_adamw_eps if kwargs_adamw_eps is not None else kwargs_eps) + if eps_value is None: + eps_value = default_eps + eps = float(eps_value) if eps_value is not None else None + + # These are handled by the shared optimizer factory and should not remain duplicated. + optimizer_kwargs.pop("fused", None) + optimizer_kwargs.pop("foreach", None) + + if optimizer_type in {"sgd", "signsgd"}: + betas = None + eps = None + + if learning_rate <= 0: + raise ValueError(f"learning_rate must be positive, got {learning_rate}") + if weight_decay < 0: + raise ValueError(f"weight_decay must be non-negative, got {weight_decay}") + if optimizer_dtype not in {"bf16", "fp32"}: + raise ValueError(f"optimizer_dtype must be 'bf16' or 'fp32', got {optimizer_dtype!r}") + + return { + "type": optimizer_type, + "learning_rate": learning_rate, + "weight_decay": weight_decay, + "optimizer_dtype": optimizer_dtype, + "betas": betas, + "eps": eps, + "optimizer_kwargs": optimizer_kwargs, + } + + +def normalize_session_spec( + *, + base_model: str, + raw_lora_config: Optional[Dict[str, Any]], + raw_optimizer_config: Optional[Dict[str, Any]], + default_rank: int, + default_alpha: int, + max_lora_rank: int, + default_optimizer_type: str, + default_learning_rate: float, + default_weight_decay: float, + default_optimizer_dtype: str, + default_optimizer_kwargs: Optional[Dict[str, Any]], + server_lora_config: Optional[Dict[str, Any]] = None, + default_betas: tuple[float, float] = DEFAULT_ADAM_BETAS, + default_eps: float = DEFAULT_ADAM_EPS, +) -> Dict[str, Any]: + """Normalize the full per-session runtime spec.""" + return { + "base_model": base_model, + "is_lora": True, + "lora_config": normalize_lora_runtime_config( + raw_lora_config, + default_rank=default_rank, + default_alpha=default_alpha, + max_lora_rank=max_lora_rank, + server_lora_config=server_lora_config, + ), + "optimizer_config": normalize_optimizer_config( + raw_optimizer_config, + default_type=default_optimizer_type, + default_learning_rate=default_learning_rate, + default_weight_decay=default_weight_decay, + default_optimizer_dtype=default_optimizer_dtype, + default_optimizer_kwargs=default_optimizer_kwargs, + default_betas=default_betas, + default_eps=default_eps, + ), + } + + +def build_default_session_spec( + *, + base_model: str, + train_config: Dict[str, Any], + lora_config: Dict[str, Any], +) -> Dict[str, Any]: + """Build the default worker session spec from server config.""" + max_lora_rank = int(lora_config.get("max_lora_rank", lora_config.get("lora_rank", 32))) + return normalize_session_spec( + base_model=base_model, + raw_lora_config={ + "lora_rank": lora_config.get("lora_rank", 32), + "lora_alpha": lora_config.get("lora_alpha", 16), + }, + raw_optimizer_config={ + "type": train_config.get("optimizer", "adamw"), + "learning_rate": train_config.get("lr", 1e-5), + "weight_decay": train_config.get("weight_decay", 0.01), + "optimizer_dtype": train_config.get("optimizer_dtype", "bf16"), + "optimizer_kwargs": train_config.get("optimizer_kwargs", {}), + }, + default_rank=lora_config.get("lora_rank", 32), + default_alpha=lora_config.get("lora_alpha", 16), + max_lora_rank=max_lora_rank, + default_optimizer_type=train_config.get("optimizer", "adamw"), + default_learning_rate=train_config.get("lr", 1e-5), + default_weight_decay=train_config.get("weight_decay", 0.01), + default_optimizer_dtype=train_config.get("optimizer_dtype", "bf16"), + default_optimizer_kwargs=train_config.get("optimizer_kwargs", {}), + server_lora_config=lora_config, + ) + + +def session_optimizer_build_kwargs(optimizer_config: Dict[str, Any]) -> Dict[str, Any]: + """Convert a normalized optimizer spec to build_optimizer kwargs.""" + kwargs = { + "lr": float(optimizer_config["learning_rate"]), + "weight_decay": float(optimizer_config.get("weight_decay", 0.0)), + "optimizer_type": optimizer_config.get("type", "adamw"), + "optimizer_dtype": optimizer_config.get("optimizer_dtype", "bf16"), + "optimizer_kwargs": _restore_optimizer_metadata(optimizer_config.get("optimizer_kwargs", {})), + } + + betas = optimizer_config.get("betas") + if betas is not None: + kwargs["betas"] = (float(betas[0]), float(betas[1])) + + eps = optimizer_config.get("eps") + if eps is not None: + kwargs["eps"] = float(eps) + + if "cautious_weight_decay" in optimizer_config: + kwargs["cautious_weight_decay"] = bool(optimizer_config["cautious_weight_decay"]) + + return kwargs + + +def write_session_spec(path: str, session_spec: Dict[str, Any]) -> str: + """Write ``session_spec.json`` to a checkpoint directory.""" + os.makedirs(path, exist_ok=True) + spec_path = os.path.join(path, SESSION_SPEC_FILENAME) + with open(spec_path, "w") as f: + json.dump(_clone_jsonable(session_spec), f, indent=2, sort_keys=True) + return spec_path + + +def read_session_spec(path: str) -> Dict[str, Any]: + """Read the normalized session spec from ``session_spec.json``.""" + spec_path = os.path.join(path, SESSION_SPEC_FILENAME) + with open(spec_path, "r") as f: + return json.load(f) + + +def session_spec_exists(path: str) -> bool: + """Return whether ``session_spec.json`` exists in a checkpoint directory.""" + return os.path.exists(os.path.join(path, SESSION_SPEC_FILENAME)) + + +def _default_betas_from_optimizer_metadata(optimizer_config: Dict[str, Any]) -> tuple[float, float]: + """Return a safe default Adam beta tuple for legacy checkpoint upgrades.""" + betas = optimizer_config.get("betas") + if betas is None: + return DEFAULT_ADAM_BETAS + return tuple(betas) + + +def load_session_spec_from_checkpoint( + path: str, + *, + fallback_base_model: Optional[str] = None, + fallback_session_spec: Optional[Dict[str, Any]] = None, +) -> Dict[str, Any]: + """Load session metadata from a checkpoint directory. + + New checkpoints write ``session_spec.json``. Older checkpoints are upgraded + from ``metadata.json`` and ``adapter_config.json`` when possible. + """ + if session_spec_exists(path): + return read_session_spec(path) + + metadata_path = os.path.join(path, "metadata.json") + adapter_config_path = os.path.join(path, "adapter_config.json") + + metadata: Dict[str, Any] = {} + adapter_config: Dict[str, Any] = {} + if os.path.exists(metadata_path): + with open(metadata_path, "r") as f: + metadata = json.load(f) + if os.path.exists(adapter_config_path): + with open(adapter_config_path, "r") as f: + adapter_config = json.load(f) + + fallback_session_spec = dict(fallback_session_spec or {}) + fallback_lora = dict(fallback_session_spec.get("lora_config") or {}) + fallback_optimizer = dict(fallback_session_spec.get("optimizer_config") or {}) + has_lora_artifacts = os.path.exists(os.path.join(path, "adapter_model.safetensors")) or bool(adapter_config) + + base_model = ( + adapter_config.get("base_model_name_or_path") or fallback_base_model or fallback_session_spec.get("base_model") + ) + if not base_model and not has_lora_artifacts: + raise FileNotFoundError( + f"Checkpoint at {path} does not contain {SESSION_SPEC_FILENAME} and no base_model fallback was provided." + ) + base_model = base_model or "" + + if not has_lora_artifacts: + return { + "base_model": base_model, + "is_lora": False, + } + + optimizer_metadata = metadata.get("optimizer") or {} + optimizer_config = normalize_optimizer_config( + { + "type": optimizer_metadata.get("type", fallback_optimizer.get("type", "adamw")), + "learning_rate": metadata.get("lr", fallback_optimizer.get("learning_rate", 1e-5)), + "weight_decay": optimizer_metadata.get("weight_decay", fallback_optimizer.get("weight_decay", 0.01)), + "optimizer_dtype": optimizer_metadata.get( + "dtype", + fallback_optimizer.get("optimizer_dtype", "bf16"), + ), + "betas": optimizer_metadata.get("betas", fallback_optimizer.get("betas")), + "eps": optimizer_metadata.get("eps", fallback_optimizer.get("eps")), + "optimizer_kwargs": optimizer_metadata.get( + "optimizer_kwargs", + fallback_optimizer.get("optimizer_kwargs", {}), + ), + }, + default_type=fallback_optimizer.get("type", "adamw"), + default_learning_rate=fallback_optimizer.get("learning_rate", 1e-5), + default_weight_decay=fallback_optimizer.get("weight_decay", 0.01), + default_optimizer_dtype=fallback_optimizer.get("optimizer_dtype", "bf16"), + default_optimizer_kwargs=fallback_optimizer.get("optimizer_kwargs", {}), + default_betas=_default_betas_from_optimizer_metadata(fallback_optimizer) + if fallback_optimizer + else DEFAULT_ADAM_BETAS, + default_eps=float(fallback_optimizer.get("eps", DEFAULT_ADAM_EPS)) + if fallback_optimizer.get("eps") is not None + else DEFAULT_ADAM_EPS, + ) + + lora_config = normalize_lora_runtime_config( + { + "lora_rank": adapter_config.get("r", fallback_lora.get("lora_rank", 32)), + "lora_alpha": adapter_config.get( + "lora_alpha", fallback_lora.get("lora_alpha", adapter_config.get("r", 32)) + ), + }, + default_rank=fallback_lora.get("lora_rank", 32), + default_alpha=fallback_lora.get("lora_alpha", 16), + max_lora_rank=max( + int(adapter_config.get("r", fallback_lora.get("lora_rank", 32))), + int(fallback_lora.get("lora_rank", 32)), + ), + ) + + return { + "base_model": base_model, + "is_lora": True, + "lora_config": lora_config, + "optimizer_config": optimizer_config, + } diff --git a/src/xorl/server/utils/zmq_channels.py b/src/xorl/server/utils/zmq_channels.py index 349d92f6..c0373d0a 100644 --- a/src/xorl/server/utils/zmq_channels.py +++ b/src/xorl/server/utils/zmq_channels.py @@ -4,9 +4,9 @@ All channels deal in raw bytes — serialization is the caller's responsibility. Channel types: -- SyncPushChannel: Sync PUSH socket (bind, send) +- SyncPushChannel: Sync PUSH socket (bind/connect, send) - SyncDealerChannel: Sync DEALER socket (connect, poll, recv) -- AsyncPullChannel: Async PULL socket (connect, poll, recv) +- AsyncPullChannel: Async PULL socket (bind/connect, poll, recv) - AsyncRouterChannel: Async ROUTER socket (bind, identity-routed send/recv) - AsyncDealerChannel: Async DEALER socket (connect, send/recv with timeouts) """ @@ -28,7 +28,7 @@ class SyncPushChannel: - """Sync PUSH socket that binds and sends single-frame messages.""" + """Sync PUSH socket that sends single-frame messages.""" def __init__(self, address: str, *, hwm: int = 1000, send_timeout: int = 1000): self._address = address @@ -47,6 +47,16 @@ def bind(self) -> None: self._socket.bind(self._address) logger.info(f"SyncPushChannel bound to {self._address}") + def connect(self) -> None: + """Create context, socket, and connect.""" + self._context = zmq.Context() + self._socket = self._context.socket(zmq.PUSH) + self._socket.setsockopt(zmq.LINGER, 0) + self._socket.setsockopt(zmq.SNDHWM, self._hwm) + self._socket.setsockopt(zmq.SNDTIMEO, self._send_timeout) + self._socket.connect(self._address) + logger.info(f"SyncPushChannel connected to {self._address}") + def send(self, data: bytes) -> None: """Send a single-frame message.""" self._socket.send(data, copy=False) @@ -117,7 +127,7 @@ def close(self) -> None: class AsyncPullChannel: - """Async PULL socket that connects and receives via polling.""" + """Async PULL socket that receives via polling.""" def __init__( self, @@ -142,6 +152,16 @@ def connect(self) -> None: self._socket.connect(self._address) logger.info(f"AsyncPullChannel connected to {self._address}") + def bind(self) -> None: + """Create socket (and context if needed) and bind.""" + if self._owns_context: + self._context = zmq.asyncio.Context() + self._socket = self._context.socket(zmq.PULL) + self._socket.setsockopt(zmq.LINGER, 0) + self._socket.setsockopt(zmq.RCVHWM, self._hwm) + self._socket.bind(self._address) + logger.info(f"AsyncPullChannel bound to {self._address}") + async def poll(self, timeout_ms: int = 100) -> bool: """Poll for incoming data. Returns True if data available.""" events = await self._socket.poll(timeout=timeout_ms) diff --git a/src/xorl/server/weight_sync/README.md b/src/xorl/server/weight_sync/README.md index c20a7cfe..a5441154 100644 --- a/src/xorl/server/weight_sync/README.md +++ b/src/xorl/server/weight_sync/README.md @@ -48,7 +48,7 @@ resp = requests.post("http://localhost:6000/api/v1/sync_inference_weights", json "master_address": "localhost", # training server address for NCCL rendezvous "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) + "flush_cache": False, # set True to flush KV cache after sync "pause_mode": "retract", # "retract" | "abort" | "in_place" # "quantization": {...} # override per-call; otherwise uses set_sync_quantization }) @@ -137,8 +137,8 @@ For each PP stage (sequential): PP stages 1+: send bf16 buffer to rank 0 via pp_group → rank 0 transfers │ ▼ +Senders: backend.complete_sync() or backend.destroy() Rank 0: endpoint_mgr.resume() -Senders: backend.destroy() All ranks: barrier ``` @@ -152,13 +152,14 @@ Key property: only one module's weights are live in GPU memory at a time ```python class WeightTransportBackend: def initialize(self) -> bool: ... # establish connections (sender ranks only) - def destroy(self) -> None: ... # tear down connections + def destroy(self, *, complete_receiver: bool = True) -> None: ... def transfer_bucket( self, bucket: List[Tuple[str, torch.Tensor]], *, src_rank: int = 0, flush_cache: bool = False, # True on the final bucket of a sync + weight_version: Optional[str] = None, ) -> None: ... # Topology hints (read by handler) @@ -209,6 +210,206 @@ Training rank 0 ──NCCL broadcast──► SGLang TP workers (ranks 1..N) - `sender_ranks = {0}` — only rank 0 sends; other training ranks only participate in training-side FSDP collectives. +## P2P Mooncake HCA Pinning + +For the P2P backend, NCCL HCA settings are not enough. Mooncake creates its own +transfer engines, so trainer ranks and SGLang receiver ranks should be pinned to +usable HCAs explicitly. + +P2P needs the Mooncake transfer engine in the trainer environment, and the +receiver must run an SGLang build with `--enable-rdma-weight-updates`. The base +`pyproject.toml` pins `mooncake-transfer-engine` so `uv sync` installs the +Python extension; the launcher image still needs CUDA runtime libraries visible +at runtime. If SGLang's `MooncakeTransferEngine` wrapper is not importable on the +trainer, xorl constructs `mooncake.engine.TransferEngine` directly. + +Trainer-side options, in precedence order: + +- `P2P_TRAINER_IB_DEVICES_PER_RANK`: semicolon-separated HCA list. If the list + covers `world_size`, entries are global-rank indexed; otherwise entries are + local-rank indexed. +- `P2P_TRAINER_GPU_TO_IB_DEVICE_MAP`: physical GPU to HCA map, for example + `0=mlx5_2,1=mlx5_3,2=mlx5_1,3=mlx5_5,4=mlx5_9,5=mlx5_9,6=mlx5_6,7=mlx5_5`. + If the launcher sets `CUDA_VISIBLE_DEVICES` to GPU UUIDs, also set + `P2P_TRAINER_VISIBLE_GPU_INDICES` to the selected physical GPU indices in + local-rank order. +- `P2P_TRAINER_IB_DEVICE`: single HCA fallback. This is useful for debugging, + but it pins every trainer rank to one rail. + +Receiver-side SGLang uses `--mooncake-ib-device` as a JSON map keyed by local +rank on each receiver node, not global TP rank. On the current H100 validation +nodes, we avoid `mlx5_4`, `mlx5_7`, and `mlx5_8` and spread TP ranks over the +remaining working HCAs. + +### Recommended P2P profile for scaled Qwen3-style MoE + +For the 4 trainer pod → 16 SGLang TP2 encoded-reasoning shape, use the +following profile as the starting point. It keeps dense/root chunking separate +from MoE batching, uses the cached receiver prepare path on warm syncs, and +avoids the measured-regressed debug/experimental knobs. + +```bash +# Required for Kubernetes Mooncake reachability. +export P2P_TRAINER_HOSTNAME="${POD_IP}" +export XORL_WEIGHT_SYNC_MASTER_ADDRESS="${POD_IP}" + +# Keep dense/root tensors small enough for scratch pools while batching MoE. +export XORL_WEIGHT_SYNC_DENSE_BUCKET_BYTES=134217728 # 128 MiB +export XORL_WEIGHT_SYNC_MOE_BUCKET_BYTES=1073741824 # 1 GiB +export XORL_WEIGHT_SYNC_BUCKET_BYTES=1073741824 # legacy MoE alias +export XORL_WEIGHT_SYNC_BATCH_MOE=1 + +# Source-reuse path keeps the required pool size near source bytes, not +# receiver-fanout bytes. 2 GiB was the best measured pool size for the scaled +# Qwen3-30B-A3B TP2 receiver layout. +export XORL_P2P_CPU_SCRATCH_POOL_BYTES=2147483648 # 2 GiB +export XORL_P2P_MOONCAKE_TRANSFER_CHUNK=8 + +# This is now the default copy mode, but keep the explicit variable in older +# generated manifests that still set XORL_P2P_SCATTER_COPY_MODE=list. +export XORL_P2P_SCATTER_REUSE_LOCATORS=1 +``` + +Leave these unset for the default performance path: + +- `XORL_P2P_USE_ASYNC_API`: Mooncake async writes are still experimental; they + have produced hangs or mixed results in repeated-update tests. +- `XORL_P2P_CPU_POOL_MIN_BYTES=0`: forces tiny transfers through CPU scratch; + this was safe in smoke tests but slower than the default GPU-direct threshold. +- `XORL_P2P_PERSIST_SMALL_REGISTRATION=1`: persistent registration of small + CUDA sources was safe in smoke tests but regressed warm sync on the scaled + TP2 layout. +- `XORL_P2P_LOG_BUCKET_DETAILS=1` and `XORL_P2P_TRANSFER_DEBUG=1`: useful for + failure diagnosis, but intentionally off the hot path because they add + logging and debug-object allocation. + +Expected warm-sync markers with this profile are +`cached_prepare=True`, `tensor_map_endpoints=0/`, near-zero +`backend_init_s`, and no SGLang receiver tensor-map payload on the second and +later syncs. + +P2P tuning options: + +- FP8 P2P sync requires an explicit sync quantization config, for example via + `POST /api/v1/set_sync_quantization` or a per-call `quantization` field: + `{"quant_method":"fp8","fmt":"e4m3","weight_block_size":[128,128]}`. + Client wrappers may expose this as `XORL_WEIGHT_SYNC_QUANTIZATION` or + `XORL_SYNC_QUANTIZATION`. A launch-only SGLang `--quantization fp8` flag is + not enough unless endpoint auto-detection is confirmed to populate the sync + request's `quantization` field. +- With P2P and explicit FP8 sync quantization, the handler quantizes supported + projection weights on the trainer side, transfers FP8 weights plus + `weight_scale_inv` tensors, and automatically asks the receiver to run + post-processing after loading. If the receiver is FP8 but the sync request has + no FP8 quantization config, tensor-size validation should fail instead of + silently copying bf16 into FP8 locators. +- The SGLang receiver must expose a matching block-FP8 layout. XORL emits + block-wise `weight_scale_inv` tensors; a receiver exposing only per-tensor + `weight_scale` tensors for FusedMoE is not compatible with this sender path. +- `XORL_P2P_FP8_QUANTIZE_DEVICE=gpu`: use the existing GPU block-FP8 kernel for + trainer-side FP8 formatting before copying the FP8 output to CPU for P2P + staging. Leave unset for the portable CPU implementation. +- `XORL_P2P_FP8_PINNED_CPU_COPY=1`: use pinned CPU output buffers for P2P FP8 + staging. This is enabled by default; set to `0` only for debugging. +- `XORL_P2P_FP8_CPU_WORKSPACE=1`: use persistent CPU workspaces for direct-EP + MoE FP8 formatting. This avoids repeated large CPU allocations and keeps the + staged HF-layout source, FP32 work buffer, abs buffer, FP8 output, and + `weight_scale_inv` output alive across syncs. +- `XORL_P2P_FP8_CPU_WORKSPACE_PINNED=1`: allocate the workspace input buffer as + pinned CPU memory when CUDA is available. Enabled by default for the workspace + path. +- `XORL_P2P_FP8_CPU_WORKSPACE_MIN_CAPACITY`: minimum expert-record capacity for + a new CPU workspace. Default: 16. +- `XORL_P2P_FP8_CPU_WORKSPACE_STREAMING=1`: stream final workspace chunks + through the P2P backend while the next chunk is being quantized. Enabled by + default for the workspace path. +- `XORL_P2P_FP8_CPU_WORKSPACE_STREAM_BYTES`: maximum quantized workspace chunk + size for streaming. Defaults to the active MoE bucket size. +- `XORL_P2P_FP8_CPU_WORKSPACE_PENDING_SOURCE_BYTES`: maximum staged BF16 source + bytes per rank before a CPU-workspace MoE batch is quantized, transferred, and + reused. Defaults to the active MoE bucket size. +- `XORL_WEIGHT_SYNC_BATCH_MOE=1`: batch direct-EP MoE expert transfers across + layers so each rank ships fewer large P2P buckets. +- The P2P backend stages each unique source tensor slice once per bucket and + reuses that pinned source address across receiver sessions. This keeps the + scratch pool sized to source bytes rather than receiver-fanout bytes. +- `XORL_P2P_BACKEND_CACHE=1`: cache P2P receiver locators and backend state + across sync calls. This is enabled by default. +- `XORL_P2P_PREPARE_WORKERS`: number of concurrent + `/prepare_weights_update` calls from the trainer to receiver endpoints. + Defaults to all endpoints, capped at 32. Set to `1` only for debugging + serialized prepare behavior. +- `XORL_P2P_PREPARE_TIMEOUT_S`: per-endpoint prepare HTTP timeout. Default: + 120 seconds. +- `XORL_P2P_SCATTER_COPY_MODE`: controls how rank 0 builds per-sender tensor + map payloads for direct-EP scatter. Default `none` reuses read-only locator + lists/dicts while constructing scatter payloads. Set `list` to shallow-copy + lists or `deep` to copy every locator dict for debugging. +- `XORL_P2P_SCATTER_REUSE_LOCATORS`: legacy boolean alias for the default + scatter copy mode. Set `1` to force locator reuse even when older manifests + still set `XORL_P2P_SCATTER_COPY_MODE=list`; set `0` to force shallow list + copies when `XORL_P2P_SCATTER_COPY_MODE` is unset. +- `XORL_WEIGHT_SYNC_MOE_BUCKET_BYTES`: explicit MoE bucket cap override. + Without this override, P2P uses a 2 GiB MoE bucket cap to amortize + Mooncake fixed costs; non-P2P backends keep the 256 MiB default. +- `XORL_WEIGHT_SYNC_BUCKET_BYTES`: legacy alias for the MoE bucket cap. Prefer + `XORL_WEIGHT_SYNC_MOE_BUCKET_BYTES` so dense/root chunking stays independent + from MoE batching. +- `XORL_P2P_USE_ASYNC_API=1`: opt into Mooncake's async write API. The default + synchronous API path is the sustained-test path; async status polling has + shown repeated-update `status=-1` failures and should remain experimental. +- `XORL_P2P_ASYNC_MIN_BYTES`: minimum coalesced chunk size for Mooncake's async + write API when `XORL_P2P_USE_ASYNC_API=1`. Default: 128 MiB. +- `XORL_P2P_MOONCAKE_WORKERS`: number of concurrent Mooncake transfer worker + calls per trainer rank. Default: 2. +- `XORL_P2P_NUM_POOLS`: number of CPU pinned scratch pools used for pipelined + staging. Default: 2. +- `XORL_P2P_MOONCAKE_TRANSFER_CHUNK`: number of coalesced staged transfers to + group into each Mooncake call. Default: 1. +- `XORL_P2P_SMALL_TRANSFER_CHUNK`: number of tiny GPU-direct transfers to group + into each Mooncake call after the per-bucket small-buffer registration. + Default: 32. +- `XORL_P2P_PERSIST_SMALL_REGISTRATION=1`: persistently register no-copy tiny + CUDA source regions across buckets/syncs. Disabled by default while being + benchmarked. +- `XORL_P2P_CPU_SCRATCH_POOL_BYTES`: CPU pinned staging pool size. Keep this + above the largest unique-source staged P2P bucket; the default is 4 GiB. +- `XORL_P2P_LOG_BUCKET_DETAILS=1`: opt into per-bucket P2P coalescing, source + reuse, and worker transfer summaries. Disabled by default to keep log I/O off + the weight-sync hot path. +- `XORL_P2P_TRANSFER_DEBUG=1`: opt into per-locator transfer debug samples in + failure messages. Disabled by default to avoid allocating debug objects for + every coalesced transfer on the hot path. +- `MC_IB_PCI_RELAXED_ORDERING=1`: enables relaxed PCIe ordering in Mooncake + RDMA when the deployment fabric supports it. Leave unset or `0` if the NIC / + platform combination is not validated. + +## Sparse Delta Probe + +`scripts/weight_sync_delta_probe.py` can measure whether an update is sparse +enough for a future sparse-delta receiver protocol to be worthwhile. It uses the +optional `delta-encoding` package when available, but it does not change the +current production P2P path. Current SGLang P2P receivers register dense tensor +buffers and expect full tensor writes; sparse deltas would also require a +receiver-side decode/scatter finalization path. + +Example: + +```bash +python scripts/weight_sync_delta_probe.py \ + --delta-encoding-path /path/to/delta-encoding \ + --shape 4096x4096 \ + --dtype uint8 \ + --density 0.001 \ + --density 0.01 \ + --density 0.1 +``` + +For dense FP8 updates, the packed sparse format is larger than the dense payload +because it stores values plus index deltas. It becomes attractive only when the +changed-entry fraction is small enough, or if a future protocol transfers LoRA +adapter tensors/factors instead of merged dense weights. + ## Adding a New Backend 1. **Create `backends/my_backend.py`** and subclass `WeightTransportBackend`: @@ -222,10 +423,11 @@ class MyBackend(WeightTransportBackend): # Use self.config.backend_config for backend-specific settings return True - def destroy(self) -> None: - # Tear down connections + def destroy(self, *, complete_receiver: bool = True) -> None: + # Tear down connections. If complete_receiver=False, skip receiver-side + # finalization because the sync failed or was aborted. - def transfer_bucket(self, bucket, *, src_rank=0, flush_cache=False): + def transfer_bucket(self, bucket, *, src_rank=0, flush_cache=False, weight_version=None): # Send [(name, tensor), ...] to inference # flush_cache=True signals the final bucket of a sync — use it to # trigger "load all weights now" for storage-based backends diff --git a/src/xorl/server/weight_sync/backends/__init__.py b/src/xorl/server/weight_sync/backends/__init__.py index dcb8df5c..b3f49502 100644 --- a/src/xorl/server/weight_sync/backends/__init__.py +++ b/src/xorl/server/weight_sync/backends/__init__.py @@ -29,4 +29,12 @@ def create_backend( from .nccl_broadcast import NCCLBroadcastBackend # noqa: PLC0415 return NCCLBroadcastBackend(config, **kwargs) - raise ValueError(f"Unknown weight sync backend: {method!r}. Supported: 'nccl_broadcast'.") + if method == "nccl_simple": + from .nccl_simple import NCCLSimpleBackend # noqa: PLC0415 + + return NCCLSimpleBackend(config, **kwargs) + if method == "p2p": + from .p2p import P2PTransportBackend # noqa: PLC0415 + + return P2PTransportBackend(config, **kwargs) + raise ValueError(f"Unknown weight sync backend: {method!r}. Supported: 'nccl_broadcast', 'nccl_simple', 'p2p'.") diff --git a/src/xorl/server/weight_sync/backends/base.py b/src/xorl/server/weight_sync/backends/base.py index 8d534f87..9f02c207 100644 --- a/src/xorl/server/weight_sync/backends/base.py +++ b/src/xorl/server/weight_sync/backends/base.py @@ -108,10 +108,16 @@ def initialize(self) -> bool: """ @abstractmethod - def destroy(self) -> None: + def destroy(self, *, complete_receiver: bool = True) -> None: """Tear down connections and free resources. Safe to call even if :meth:`initialize` was not called or failed. + + Args: + complete_receiver: If ``False``, skip any receiver-side "sync + complete" action and only clean up local transport resources. + Failure/abort paths use this to avoid marking a partial sync as + complete. """ # ------------------------------------------------------------------ @@ -142,6 +148,20 @@ def transfer_bucket( this on the final bucket of a sync. """ + def flush_pending_transfers(self) -> None: + """Block until any in-flight async transfers complete. + + For backends that issue ``transfer_bucket`` synchronously (e.g. + the NCCL broadcaster), this is a no-op — by the time + ``transfer_bucket`` returns, the bytes have landed. + + For async backends (P2P/Mooncake), bucket calls return after + staging and submitting work to a worker thread; the handler + must call this before resuming inference, otherwise generation + can resume on partially-updated weights. + """ + return None + # ------------------------------------------------------------------ # Topology hints (read by the handler to decide who prepares data) # ------------------------------------------------------------------ diff --git a/src/xorl/server/weight_sync/backends/nccl_broadcast.py b/src/xorl/server/weight_sync/backends/nccl_broadcast.py index ec4f3f65..950ec1b6 100644 --- a/src/xorl/server/weight_sync/backends/nccl_broadcast.py +++ b/src/xorl/server/weight_sync/backends/nccl_broadcast.py @@ -55,6 +55,19 @@ def _get_http_session() -> requests.Session: return _http_session +def _ws_port(endpoint: "EndpointInfo") -> int: + """Weight-sync receiver port for an endpoint. + + XORL_WEIGHT_SYNC_PORT overrides the endpoint's serving port so the + init/update/destroy HTTP calls can target a sidecar receiver process + (which receives the NCCL broadcast and hands tensors to the inference + server via /update_weights_from_tensor) while generation traffic and + pause/resume keep using the real serving port. + """ + override = os.environ.get("XORL_WEIGHT_SYNC_PORT") + return int(override) if override else endpoint.port + + @dataclass class EndpointInfo: """Information about an inference endpoint.""" @@ -155,12 +168,20 @@ def _create_training_store(self) -> None: logger.info(f"[Training] Creating TCPStore (requested_port={requested_port}, is_master=True)...") with self._without_torchelastic_agent_store(): + # wait_for_workers=False: master must start listening without blocking. + # /init_weights_update_group is only sent to inference endpoints later, + # in init_inference (started after this call returns), and the actual + # NCCL rendezvous is completed inside _init_training_process_group via + # _new_process_group_helper. With the default wait_for_workers=True, + # construction blocks waiting for workers that cannot connect yet, + # deadlocking sync_inference_weights. raw_store = TCPStore( host_name=self.master_address, port=requested_port, world_size=self.world_size, is_master=True, timeout=default_pg_timeout, + wait_for_workers=False, ) self._training_raw_store = raw_store @@ -228,7 +249,10 @@ def _init_training_process_group(self) -> dist.ProcessGroup: else: pg_options_param_name = "pg_options" - # Set CUDA device and pass device_id for proper NCCL comm initialization + # Use eager NCCL init (device_id). This works correctly for 2-node + # cross-node setup (IB network, no CUDA visibility issues). + # Was broken on same-node split CUDA_VISIBLE_DEVICES due to sglang's + # internal TP broadcast deadlock; 2-node avoids that entirely. torch.cuda.set_device(self.device) device_id = torch.device(self.device) @@ -358,9 +382,19 @@ def destroy_nccl_group(self) -> None: if self.process_group: try: - dist.destroy_process_group(self.process_group) + # Use the non-cooperative abort() rather than + # dist.destroy_process_group(), which calls pg.shutdown() and + # then waits for the inference-side comm to finalize. The + # inference side aborted its comm immediately on receiving + # /destroy_weights_update_group, so a cooperative shutdown + # here hangs until the engine timeout fires. + self.process_group.abort() except Exception as e: - logger.error(f"Failed to destroy training process group: {e}") + logger.warning(f"process_group.abort() failed ({e}); falling back to destroy_process_group") + try: + dist.destroy_process_group(self.process_group) + except Exception as e2: + logger.error(f"Failed to destroy training process group: {e2}") self.process_group = None self._cleanup_training_store() @@ -381,7 +415,7 @@ def _init_inference_endpoints(self) -> List[Dict[str, Any]]: session = _get_http_session() def init_single(rank_offset: int, endpoint: EndpointInfo) -> Dict[str, Any]: - url = f"http://{endpoint.host}:{endpoint.port}/init_weights_update_group" + url = f"http://{endpoint.host}:{_ws_port(endpoint)}/init_weights_update_group" payload = { "master_address": self.master_address, "master_port": self._active_master_port, @@ -437,7 +471,7 @@ def _destroy_inference_endpoints(self) -> List[Dict[str, Any]]: session = _get_http_session() def destroy_single(endpoint: EndpointInfo) -> Dict[str, Any]: - url = f"http://{endpoint.host}:{endpoint.port}/destroy_weights_update_group" + url = f"http://{endpoint.host}:{_ws_port(endpoint)}/destroy_weights_update_group" payload = {"group_name": self.group_name} try: response = session.post(url, json=payload, timeout=30) @@ -522,7 +556,7 @@ def _endpoint_request_with_retry( def pause_inference_endpoints( self, - pause_mode: str = "retract", + pause_mode: str = "in_place", max_retries: int = 3, retry_delay_seconds: float = 1.0, ) -> Tuple[List[Dict[str, Any]], bool]: @@ -632,8 +666,9 @@ def _transfer_single_bucket( def call_single_endpoint(endpoint: EndpointInfo, endpoint_idx: int): """Call update_weights_from_distributed on a single endpoint.""" try: + logger.info(f"[Training] HTTP thread: posting update_weights to {endpoint.host}:{endpoint.port}") response = session.post( - f"http://{endpoint.host}:{endpoint.port}/update_weights_from_distributed", + f"http://{endpoint.host}:{_ws_port(endpoint)}/update_weights_from_distributed", json={ "names": names, "dtypes": dtypes, @@ -655,6 +690,7 @@ def call_single_endpoint(endpoint: EndpointInfo, endpoint_idx: int): if not result.get("success"): update_errors.append(f"API failed on {endpoint.host}:{endpoint.port}: {result}") except Exception as e: + logger.error(f"[Training] HTTP thread failed for {endpoint.host}:{endpoint.port}: {e}") update_errors.append(f"Exception calling {endpoint.host}:{endpoint.port}: {e}") # Start API calls in parallel threads (one per endpoint) @@ -664,6 +700,15 @@ def call_single_endpoint(endpoint: EndpointInfo, endpoint_idx: int): t.start() api_threads.append(t) + # Fail fast if the HTTP request errored immediately (connection + # refused etc.). Without this check, dist.broadcast below blocks + # forever waiting for receivers that were never notified. + time.sleep(0.5) + if update_errors: + for t in api_threads: + t.join(timeout=5) + raise RuntimeError(f"update_weights HTTP failed before broadcast: {update_errors}") + # Set device once for all operations torch.cuda.set_device(self.device) @@ -685,6 +730,16 @@ def call_single_endpoint(endpoint: EndpointInfo, endpoint_idx: int): dist.broadcast(param_data, src=0, group=self.process_group) + # Force CUDA stream to finish so receiver-side handle.wait() can + # actually progress; without this, dist.broadcast can return at the + # Python level while the NCCL kernel is still queued on the stream, + # which leaves sglang blocked in update_weights_from_distributed and + # the api_threads below hung waiting for the 200 OK. + if (isinstance(self.device, torch.device) and self.device.type == "cuda") or ( + isinstance(self.device, str) and self.device.startswith("cuda") + ): + torch.cuda.synchronize(self.device) + # Wait for all API calls to complete for t in api_threads: t.join() @@ -882,7 +937,8 @@ def initialize(self) -> bool: self._process_group = self._synchronizer.process_group return ok - def destroy(self) -> None: + def destroy(self, *, complete_receiver: bool = True) -> None: + _ = complete_receiver if self._synchronizer is not None: try: self._synchronizer.destroy_nccl_group() diff --git a/src/xorl/server/weight_sync/backends/nccl_simple.py b/src/xorl/server/weight_sync/backends/nccl_simple.py new file mode 100644 index 00000000..35da4850 --- /dev/null +++ b/src/xorl/server/weight_sync/backends/nccl_simple.py @@ -0,0 +1,123 @@ +"""Two-phase NCCL weight sync backend. + +Why this exists +--------------- +``nccl_broadcast`` interleaves two NCCL communicators per module: + + unshard (FSDP all-gather, intra-node) -> extract -> reshard + -> dist.broadcast (weight-sync group, inter-node) -> next module + +With FSDP ranks 1..N racing ahead of rank 0 (they have no broadcast work), +the two communicators enqueue kernels in different orders across ranks, +which deadlocks NCCL after a few modules (observed consistently at the +7th bucket on 14B / FSDP=4 / 2-node). + +This backend removes the interleaving entirely: + +* **Phase A** (during the handler's module loop): ``transfer_bucket`` only + stages tensors to CPU. No NCCL, no HTTP. The FSDP loop runs to completion + using only FSDP collectives. +* **Phase B** (``flush_pending_transfers``, called by the handler after the + module loop): re-chunk staged params and send each chunk through the + proven ``_transfer_single_bucket`` path (HTTP + dist.broadcast). Only the + weight-sync communicator is active. + +The two communicators never run concurrently, so the kernel-ordering +deadlock cannot occur. +""" + +import logging +from typing import List, Optional, Tuple + +import torch + +from .nccl_broadcast import NCCLBroadcastBackend + + +logger = logging.getLogger(__name__) + +# Re-chunk size for phase B. Bounds sglang-side temp memory +# (torch.empty per param before load_weights) and trainer-side H2D staging. +_CHUNK_BYTES = 1024 * 1024 * 1024 # 1 GiB + + +class NCCLSimpleBackend(NCCLBroadcastBackend): + """Two-phase (stage-then-broadcast) NCCL transport.""" + + def __init__(self, config, **kwargs) -> None: + super().__init__(config, **kwargs) + self._pending: List[Tuple[str, torch.Tensor]] = [] + self._pending_bytes = 0 + self._final_flush_cache = False + self._final_weight_version: Optional[str] = None + + # ------------------------------------------------------------------ + # Phase A: stage to CPU (no NCCL, no HTTP) + # ------------------------------------------------------------------ + def transfer_bucket( + self, + bucket: List[Tuple[str, torch.Tensor]], + *, + src_rank: int = 0, + flush_cache: bool = False, + weight_version: Optional[str] = None, + ) -> None: + if src_rank != 0: + raise ValueError(f"NCCLSimpleBackend only supports src_rank=0, got {src_rank}") + for name, t in bucket: + cpu_t = t.detach().to("cpu").contiguous() + self._pending.append((name, cpu_t)) + self._pending_bytes += cpu_t.numel() * cpu_t.element_size() + if flush_cache: + self._final_flush_cache = True + if weight_version is not None: + self._final_weight_version = weight_version + + # ------------------------------------------------------------------ + # Phase B: chunked HTTP + broadcast (no FSDP collectives anywhere) + # ------------------------------------------------------------------ + def flush_pending_transfers(self) -> None: + if not self._pending: + return + if self._synchronizer is None: + raise RuntimeError("Backend not initialized — call initialize() first") + + # Build chunks bounded by _CHUNK_BYTES (a single oversized param + # becomes its own chunk). + chunks: List[List[Tuple[str, torch.Tensor]]] = [] + cur: List[Tuple[str, torch.Tensor]] = [] + cur_bytes = 0 + for name, t in self._pending: + nbytes = t.numel() * t.element_size() + if cur and cur_bytes + nbytes > _CHUNK_BYTES: + chunks.append(cur) + cur, cur_bytes = [], 0 + cur.append((name, t)) + cur_bytes += nbytes + if cur: + chunks.append(cur) + + total_gb = self._pending_bytes / 1e9 + logger.info( + f"[NCCLSimple] Phase B: broadcasting {len(self._pending)} params " + f"({total_gb:.2f} GB) in {len(chunks)} chunks" + ) + + device = self.config.device + try: + for i, chunk in enumerate(chunks): + last = i == len(chunks) - 1 + gpu_chunk = [(n, t.to(device, non_blocking=True)) for n, t in chunk] + torch.cuda.synchronize(device) + self._synchronizer._transfer_single_bucket( + gpu_chunk, + flush_cache=self._final_flush_cache and last, + weight_version=self._final_weight_version if last else None, + ) + del gpu_chunk + logger.info(f"[NCCLSimple] Phase B complete: {len(chunks)} chunks sent") + finally: + self._pending = [] + self._pending_bytes = 0 + self._final_flush_cache = False + self._final_weight_version = None diff --git a/src/xorl/server/weight_sync/backends/p2p.py b/src/xorl/server/weight_sync/backends/p2p.py new file mode 100644 index 00000000..43a32186 --- /dev/null +++ b/src/xorl/server/weight_sync/backends/p2p.py @@ -0,0 +1,2572 @@ +"""P2P (Mooncake) weight transport backend. + +RDMA one-sided writes from training ranks directly into the inference +replica's ``param.data`` slices, modeled on the lmsys/Mooncake P2P weight +update mechanism (https://www.lmsys.org/blog/2026-04-29-p2p-update/). + +Compared to the NCCL-broadcast backend: + +* No NCCL group rendezvous. +* No rank-0 dist.broadcast bottleneck. +* Each inference TP rank registers its own param memory with Mooncake; the + trainer issues writes against the per-rank session ids returned by + SGLang's ``/prepare_weights_update``. + +The backend supports the rank-0 dense path and the direct-EP MoE path. PP +stage leaders still funnel through rank 0 in the handler. +""" + +import dataclasses +import ipaddress +import logging +import os +import socket +import time +import zlib +from concurrent.futures import Future, ThreadPoolExecutor, as_completed +from typing import Any, Dict, FrozenSet, List, Optional, Tuple + +import requests +import torch +import torch.distributed as dist + +from .base import TransportConfig, WeightTransportBackend + + +@dataclasses.dataclass +class _PendingTransfer: + """One per-locator transfer pending a stage+coalesce pass in + ``transfer_bucket``. Held just long enough to sort by peer_ptr, copy + the src_view into the CPU pinned scratch pool, and coalesce + contiguous neighbors before issuing the Mooncake call. + """ + + peer_ptr: int + nbytes: int + src_view: torch.Tensor + name: str + loc: Dict[str, Any] + + +@dataclasses.dataclass +class _TransferDebugEntry: + name: str + peer_ptr: int + nbytes: int + dtype: Optional[str] + memory_handle: Optional[int] + tp_rank: Any + ep_rank: Any + loc_slice: Any + + +@dataclasses.dataclass +class _StagedTransfer: + src_ptr: int + peer_ptr: int + nbytes: int + memory_handle: Optional[int] + name: str + loc: Dict[str, Any] + + +@dataclasses.dataclass +class _TransferDebugSample: + entries: List[_TransferDebugEntry] = dataclasses.field(default_factory=list) + total: int = 0 + + def add(self, name: str, loc: Dict[str, Any], nbytes: int) -> None: + self.total += 1 + if len(self.entries) < _TRANSFER_DEBUG_SAMPLE_LIMIT: + self.entries.append(_transfer_debug_entry(name, loc, nbytes)) + + def extend(self, other: "_TransferDebugSample") -> None: + self.total += other.total + if len(self.entries) >= _TRANSFER_DEBUG_SAMPLE_LIMIT: + return + remaining = _TRANSFER_DEBUG_SAMPLE_LIMIT - len(self.entries) + self.entries.extend(other.entries[:remaining]) + + +@dataclasses.dataclass +class _BucketTiming: + """Per-bucket wall-time breakdown for the P2P transport.""" + + nbytes: int = 0 + prepare_s: float = 0.0 + pool_init_s: float = 0.0 + pool_wait_s: float = 0.0 + stage_s: float = 0.0 + submit_s: float = 0.0 + register_s: float = 0.0 + transfer_s: float = 0.0 + deregister_s: float = 0.0 + num_large_buffers: int = 0 + num_small_buffers: int = 0 + session_bytes: Dict[str, int] = dataclasses.field(default_factory=dict) + session_transfer_s: Dict[str, float] = dataclasses.field(default_factory=dict) + + @property + def total_s(self) -> float: + return ( + self.prepare_s + + self.pool_init_s + + self.pool_wait_s + + self.stage_s + + self.submit_s + + self.register_s + + self.transfer_s + + self.deregister_s + ) + + @property + def main_thread_s(self) -> float: + return self.prepare_s + self.pool_init_s + self.pool_wait_s + self.stage_s + self.submit_s + + @property + def throughput_mb_s(self) -> float: + if self.total_s <= 0: + return 0.0 + return (self.nbytes / 1e6) / self.total_s + + +logger = logging.getLogger(__name__) + + +_HTTP_TIMEOUT_SECONDS = 600 +_TRANSFER_DEBUG_SAMPLE_LIMIT = 6 + + +def _env_float(name: str, default: float) -> float: + raw = os.environ.get(name) + if raw is None: + return default + try: + value = float(raw) + except ValueError: + logger.warning("[P2P] invalid %s=%r; using %.1fs", name, raw, default) + return default + if value <= 0: + logger.warning("[P2P] invalid %s=%r; using %.1fs", name, raw, default) + return default + return value + + +def _env_int(name: str, default: int, *, minimum: int = 1) -> int: + raw = os.environ.get(name) + if raw is None: + return default + try: + value = int(raw) + except ValueError: + logger.warning("[P2P] invalid %s=%r; using %d", name, raw, default) + return default + if value < minimum: + logger.warning("[P2P] invalid %s=%r; using %d", name, raw, default) + return default + return value + + +def _env_flag(name: str, default: bool = False) -> bool: + raw = os.environ.get(name) + if raw is None: + return default + value = raw.strip().lower() + if value in {"1", "true", "yes", "on"}: + return True + if value in {"0", "false", "no", "off"}: + return False + logger.warning("[P2P] invalid %s=%r; using %s", name, raw, default) + return default + + +def _prepare_timeout_seconds() -> float: + return _env_float("XORL_P2P_PREPARE_TIMEOUT_S", 120.0) + + +def _async_api_min_bytes() -> int: + return _env_int("XORL_P2P_ASYNC_MIN_BYTES", 128 * 1024 * 1024) + + +def _small_transfer_chunk() -> int: + return _env_int("XORL_P2P_SMALL_TRANSFER_CHUNK", 32) + + +def _persist_small_registration_enabled() -> bool: + # Opt-in only. On the scaled TP2 layout this was safe, but slower than + # per-bucket small-source registration because it increased warm-sync tails. + return _env_flag("XORL_P2P_PERSIST_SMALL_REGISTRATION", False) + + +def _async_api_enabled(*, cached_prepare: bool) -> bool: + # The synchronous Mooncake API is the measured sustained path. The async + # API is kept for experiments because repeated-update tests have shown + # mixed results and hangs/status failures on some runs. + mode = os.environ.get("XORL_P2P_USE_ASYNC_API", "0").strip().lower() + if mode in {"1", "true", "yes", "on"}: + return True + if mode in {"warm", "cached", "cached_prepare"}: + return cached_prepare + if mode in {"0", "false", "no", "off"}: + return False + logger.warning("[P2P] invalid XORL_P2P_USE_ASYNC_API=%r; using sync transfer API", mode) + return False + + +def _align_up(value: int, alignment: int) -> int: + if alignment <= 1: + return value + return ((value + alignment - 1) // alignment) * alignment + + +# CPU-pinned scratch pool size for source staging. Set big +# enough to hold a full bucket's worth of source bytes — the default +# bucket cap is 2 GB, so we size 4 GB by default for safety. The pool +# is registered with Mooncake once at first use and reused for every +# bucket with no per-bucket register cost. Tunable via env var on +# memory-constrained deployments. +_CPU_SCRATCH_POOL_BYTES = int(os.environ.get("XORL_P2P_CPU_SCRATCH_POOL_BYTES", str(4 * 1024 * 1024 * 1024))) + +# Mooncake's CPU-source RDMA path on our cluster fails (ret=-1) for +# very small transfers — observed at 8 KB on a layernorm weight. Small +# entries take the GPU-direct path (per-bucket register/dereg); large +# entries take the CPU pool path. 64 KB threshold matches typical +# layernorm weight size (2048 BF16 = 4 KB; 4× headroom). Setting this +# to 0 forces tiny entries through CPU scratch; that was safe in smoke +# tests, but slower than the default GPU-direct threshold. +_CPU_POOL_MIN_BYTES = int(os.environ.get("XORL_P2P_CPU_POOL_MIN_BYTES", str(64 * 1024))) + + +class _DirectMooncakeTransferEngine: + """Minimal xorl-side wrapper for mooncake.engine.TransferEngine. + + SGLang ships a convenience wrapper with the same surface, but the trainer + environment should only need ``mooncake-transfer-engine`` to construct a + sender. Keep this fallback small and aligned with the methods used below. + """ + + def __init__( + self, + transfer_engine_cls: Any, + *, + hostname: str, + gpu_id: int, + ib_device: Optional[str], + ) -> None: + self.engine = transfer_engine_cls() + self.hostname = hostname + self.gpu_id = gpu_id + self.ib_device = ib_device + ret = self.engine.initialize( + hostname, + "P2PHANDSHAKE", + "rdma", + ib_device or "", + ) + if ret != 0: + raise RuntimeError(f"Mooncake TransferEngine initialization failed: ret={ret}") + self.session_id = f"{hostname}:{self.engine.get_rpc_port()}" + + def get_session_id(self) -> str: + return self.session_id + + def get_ib_device(self) -> Optional[str]: + return self.ib_device + + def batch_register(self, ptrs: List[int], lengths: List[int]) -> int: + try: + return self.engine.batch_register_memory(ptrs, lengths) + except Exception: + if not hasattr(self.engine, "batch_register_memory"): + raise RuntimeError("Mooncake batch register requires a newer mooncake-transfer-engine") + return -1 + + def batch_deregister(self, ptrs: List[int]) -> int: + try: + return self.engine.batch_unregister_memory(ptrs) + except Exception: + return -1 + + def batch_transfer_sync( + self, + session_id: str, + buffers: List[int], + peer_buffer_addresses: List[int], + lengths: List[int], + ) -> int: + try: + return self.engine.batch_transfer_sync_write(session_id, buffers, peer_buffer_addresses, lengths) + except Exception: + if not hasattr(self.engine, "batch_transfer_sync_write"): + raise RuntimeError("Mooncake batch transfer requires a newer mooncake-transfer-engine") + return -1 + + +class _CompletedCudaEvent: + """CPU-test stand-in for ``torch.cuda.Event``.""" + + def synchronize(self) -> None: + return None + + +def _retry_delay(attempt: int) -> float: + return min(1.0, 0.05 * (2 ** min(attempt, 4))) + + +def _locator_memory_handle(loc: Dict[str, Any]) -> Optional[int]: + raw = loc.get("memory_handle") + if raw is None: + return None + try: + return int(raw) + except (TypeError, ValueError): + return None + + +def _transfer_debug_entry(name: str, loc: Dict[str, Any], nbytes: int) -> _TransferDebugEntry: + return _TransferDebugEntry( + name=name, + peer_ptr=int(loc.get("ptr", 0)), + nbytes=nbytes, + dtype=loc.get("dtype"), + memory_handle=_locator_memory_handle(loc), + tp_rank=loc.get("tp_rank"), + ep_rank=loc.get("ep_rank"), + loc_slice=loc.get("slice"), + ) + + +def _source_view_key(src_view: torch.Tensor, nbytes: int) -> Tuple[Any, ...]: + ptr = int(src_view.data_ptr()) + nbytes = int(nbytes) + if src_view.is_contiguous(): + # For contiguous views, the byte range fully identifies the payload we + # copy into the CPU pool. Avoid tupleizing shape/stride for the common + # fanout case where many receivers share the same source slice. + return (ptr, nbytes) + return ( + ptr, + tuple(int(dim) for dim in src_view.shape), + tuple(int(stride) for stride in src_view.stride()), + str(src_view.dtype), + nbytes, + ) + + +def _format_transfer_debug(debug_entries: Any) -> str: + if debug_entries is None: + return "transfer_debug=disabled (set XORL_P2P_TRANSFER_DEBUG=1)" + + total_entries: Optional[int] = None + if isinstance(debug_entries, _TransferDebugSample): + total_entries = debug_entries.total + debug_entries = debug_entries.entries + else: + total_entries = len(debug_entries) + + if not debug_entries: + return "transfer_debug=[]" + + parts: List[str] = [] + for entry in debug_entries[:_TRANSFER_DEBUG_SAMPLE_LIMIT]: + handle = f"0x{entry.memory_handle:x}" if entry.memory_handle is not None else "None" + parts.append( + f"{entry.name}(ptr=0x{entry.peer_ptr:x}, nbytes={entry.nbytes}, " + f"dtype={entry.dtype}, handle={handle}, tp={entry.tp_rank}, " + f"ep={entry.ep_rank}, slice={entry.loc_slice})" + ) + if total_entries > len(debug_entries): + parts.append(f"... {total_entries - len(debug_entries)} more") + return "transfer_debug=[" + "; ".join(parts) + "]" + + +def _chunk_sizes( + by_session: Dict[str, Tuple[List[int], List[int], List[int], Optional[List[_TransferDebugSample]]]], + session_id: str, + i: int, + end: int, +) -> List[int]: + return by_session[session_id][2][i:end] + + +def _chunk_debug_sample( + by_session: Dict[str, Tuple[List[int], List[int], List[int], Optional[List[_TransferDebugSample]]]], + session_id: str, + i: int, + end: int, +) -> Optional[_TransferDebugSample]: + debug_entries = by_session[session_id][3] + if debug_entries is None: + return None + debug_lists = debug_entries[i:end] + sample = _TransferDebugSample() + for debug in debug_lists: + sample.extend(debug) + return sample + + +def _run_sync_transfer_items( + *, + engine_wrapper: Any, + by_session: Dict[str, Tuple[List[int], List[int], List[int], Optional[List[_TransferDebugSample]]]], + items: List[Tuple[str, int, int]], + session_debug_info: Dict[str, Dict[str, Any]], + session_transfer_s: Dict[str, float], + bucket_idx: int, + label: str, +) -> None: + max_attempts = _env_int("XORL_P2P_TRANSFER_RETRIES", 10) + for session_id, i, end in items: + src_ptrs, peer_ptrs, lengths, _ = by_session[session_id] + t_session = time.perf_counter() + last_ret = 0 + for attempt in range(max_attempts): + last_ret = engine_wrapper.batch_transfer_sync(session_id, src_ptrs[i:end], peer_ptrs[i:end], lengths[i:end]) + if last_ret >= 0: + break + time.sleep(_retry_delay(attempt)) + if last_ret < 0: + raise RuntimeError( + f"[P2P] {label} to {session_id} failed: ret={last_ret} " + f"(bucket {bucket_idx}, chunk {i}..{end} of {len(src_ptrs)} buffers, " + f"sizes={_chunk_sizes(by_session, session_id, i, end)}, " + f"{_format_transfer_debug(_chunk_debug_sample(by_session, session_id, i, end))}, " + f"session_info={session_debug_info.get(session_id)}, after {max_attempts} attempts)" + ) + session_transfer_s[session_id] = session_transfer_s.get(session_id, 0.0) + (time.perf_counter() - t_session) + + +def _run_async_transfer_items( + *, + engine_wrapper: Any, + by_session: Dict[str, Tuple[List[int], List[int], List[int], Optional[List[_TransferDebugSample]]]], + items: List[Tuple[str, int, int]], + session_debug_info: Dict[str, Dict[str, Any]], + session_transfer_s: Dict[str, float], + bucket_idx: int, +) -> None: + if not items: + return + + # Bounded submit/poll. Mooncake's underlying TransferEngine exposes: + # batch_transfer_async_write(session, src, dst, lens) -> batch_id (int) + # get_batch_transfer_status([batch_id, ...]) -> int (0 success, -1 failure/timeout) + # + # Status takes a sequence of batch IDs. The single-transfer status API is + # not valid for these batch IDs and can wedge the caller. + raw_engine = engine_wrapper.engine + max_in_flight = _env_int("XORL_P2P_ASYNC_MAX_IN_FLIGHT", 1) + status_timeout_s = max(0.001, _env_float("XORL_P2P_ASYNC_STATUS_TIMEOUT_S", 30.0)) + work_items = list(items) + active: List[Tuple[int, str, int, int, float]] = [] + active_since: Optional[float] = None + last_status_log_at = 0.0 + + while work_items or active: + while work_items and len(active) < max_in_flight: + session_id, i, end = work_items.pop(0) + src_ptrs, peer_ptrs, lengths, _ = by_session[session_id] + bid = raw_engine.batch_transfer_async_write(session_id, src_ptrs[i:end], peer_ptrs[i:end], lengths[i:end]) + if bid is None or (isinstance(bid, int) and bid < 0): + raise RuntimeError( + f"[P2P] batch_transfer_async_write submit failed: bid={bid} " + f"(bucket {bucket_idx}, chunk {i}..{end}, " + f"sizes={_chunk_sizes(by_session, session_id, i, end)}, " + f"{_format_transfer_debug(_chunk_debug_sample(by_session, session_id, i, end))}, " + f"session_info={session_debug_info.get(session_id)})" + ) + if not active: + active_since = time.perf_counter() + active.append((int(bid), session_id, i, end, time.perf_counter())) + + if not active: + active_since = None + continue + + bids = [bid for bid, *_ in active] + status = raw_engine.get_batch_transfer_status(bids) + if status < 0: + _, session_id, i, end, _ = active[0] + raise RuntimeError( + f"[P2P] get_batch_transfer_status reported failure: status={status} " + f"(bucket {bucket_idx}, {len(active)} batches in flight, " + f"first session={session_id}, sizes={_chunk_sizes(by_session, session_id, i, end)}, " + f"{_format_transfer_debug(_chunk_debug_sample(by_session, session_id, i, end))}, " + f"session_info={session_debug_info.get(session_id)})" + ) + if status == 0: + now = time.perf_counter() + for _, session_id, _, _, submit_t in active: + session_transfer_s[session_id] = session_transfer_s.get(session_id, 0.0) + (now - submit_t) + active = [] + active_since = None + continue + + now = time.perf_counter() + waited_s = now - (active_since or now) + if waited_s > status_timeout_s: + _, session_id, i, end, _ = active[0] + raise RuntimeError( + f"[P2P] async transfer status poll timed out: status={status} " + f"(bucket {bucket_idx}, waited={waited_s:.3f}s, " + f"{len(active)} batches in flight, first session={session_id}, " + f"sizes={_chunk_sizes(by_session, session_id, i, end)}, " + f"{_format_transfer_debug(_chunk_debug_sample(by_session, session_id, i, end))}, " + f"session_info={session_debug_info.get(session_id)})" + ) + if now - last_status_log_at > 5.0: + logger.warning( + f"[P2P] async transfer still pending: status={status} " + f"(bucket {bucket_idx}, waited={waited_s:.3f}s, {len(active)} batches in flight)" + ) + last_status_log_at = now + time.sleep(0.0001) + + +def _transfer_small_entries( + *, + engine_wrapper: Any, + small_session_data: Dict[str, List[Tuple[int, int, int, Optional[_TransferDebugEntry]]]], + session_debug_info: Dict[str, Dict[str, Any]], + small_register_ptrs: List[int], + small_register_lens: List[int], + session_bytes: Dict[str, int], + session_transfer_s: Dict[str, float], + bucket_idx: int, +) -> Tuple[int, int]: + if small_register_ptrs: + ret = engine_wrapper.batch_register(small_register_ptrs, small_register_lens) + if ret != 0: + raise RuntimeError(f"[P2P] small-entries batch_register failed: ret={ret} (bucket {bucket_idx})") + + total_bytes = 0 + num_buffers = 0 + try: + chunk = _small_transfer_chunk() + max_attempts = _env_int("XORL_P2P_TRANSFER_RETRIES", 10) + for session_id, triples in small_session_data.items(): + t_session = time.perf_counter() + for i in range(0, len(triples), chunk): + transfer_chunk = triples[i : i + chunk] + src_ptrs = [src_ptr for src_ptr, _, _, _ in transfer_chunk] + peer_ptrs = [peer_ptr for _, peer_ptr, _, _ in transfer_chunk] + lengths = [nbytes for _, _, nbytes, _ in transfer_chunk] + chunk_bytes = sum(lengths) + total_bytes += chunk_bytes + session_bytes[session_id] = session_bytes.get(session_id, 0) + chunk_bytes + num_buffers += len(transfer_chunk) + last_ret = 0 + for attempt in range(max_attempts): + last_ret = engine_wrapper.batch_transfer_sync(session_id, src_ptrs, peer_ptrs, lengths) + if last_ret >= 0: + break + time.sleep(_retry_delay(attempt)) + if last_ret < 0: + debug_entries = [debug for _, _, _, debug in transfer_chunk if debug is not None] + raise RuntimeError( + f"[P2P] small-entries transfer to {session_id} " + f"failed: ret={last_ret} (bucket {bucket_idx}, " + f"chunk {i}..{i + len(transfer_chunk)} of {len(triples)} buffers, " + f"sizes={lengths}, " + f"{_format_transfer_debug(debug_entries or None)}, " + f"session_info={session_debug_info.get(session_id)}, " + f"after {max_attempts} attempts)" + ) + session_transfer_s[session_id] = session_transfer_s.get(session_id, 0.0) + (time.perf_counter() - t_session) + finally: + if small_register_ptrs: + try: + engine_wrapper.batch_deregister(small_register_ptrs) + except Exception as e: + logger.warning(f"[P2P] small-entries dereg failed (bucket {bucket_idx}): {e}") + + return total_bytes, num_buffers + + +def _do_async_transfer( + *, + engine_wrapper: Any, + copy_done_event: "torch.cuda.Event", + by_session: Dict[str, Tuple[List[int], List[int], List[int], Optional[List[_TransferDebugSample]]]], + small_session_data: Dict[str, List[Tuple[int, int, int, Optional[_TransferDebugEntry]]]], + session_debug_info: Dict[str, Dict[str, Any]], + small_register_ptrs: List[int], + small_register_lens: List[int], + chunk: int, + use_async_api: bool, + timing: _BucketTiming, + bucket_idx: int, + slice_holds: List[torch.Tensor], + src_view_holds: List[torch.Tensor], + log_bucket_details: bool, +) -> None: + """Worker-thread Mooncake transfer for one bucket. + + The worker waits for CUDA staging to finish, ships large entries from the + registered CPU pool, and then handles tiny GPU-direct entries that are too + small for the CPU-source path. ``slice_holds`` and ``src_view_holds`` keep + source memory alive while the caller reshards the FSDP module. + """ + copy_done_event.synchronize() + + t0 = time.perf_counter() + bucket_bytes = 0 + session_bytes: Dict[str, int] = {} + session_transfer_s: Dict[str, float] = {} + num_large_buffers = 0 + + all_large_items: List[Tuple[str, int, int]] = [] + async_items: List[Tuple[str, int, int]] = [] + sync_fallback_items: List[Tuple[str, int, int]] = [] + async_min_bytes = _async_api_min_bytes() + for session_id, (src_ptrs, _, lengths, _) in by_session.items(): + nbytes = sum(lengths) + bucket_bytes += nbytes + session_bytes[session_id] = session_bytes.get(session_id, 0) + nbytes + num_large_buffers += len(lengths) + for i in range(0, len(src_ptrs), chunk): + end = min(i + chunk, len(src_ptrs)) + item = (session_id, i, end) + all_large_items.append(item) + if sum(lengths[i:end]) < async_min_bytes: + sync_fallback_items.append(item) + else: + async_items.append(item) + + # Mooncake's CPU-source sync path can return ret=-1 under high load, so + # sync transfers keep bounded retries. The async API is stricter: if submit + # or status fails, we fail closed because a prior async batch may still be + # writing receiver memory. + if use_async_api: + _run_sync_transfer_items( + engine_wrapper=engine_wrapper, + by_session=by_session, + items=sync_fallback_items, + session_debug_info=session_debug_info, + session_transfer_s=session_transfer_s, + bucket_idx=bucket_idx, + label="async sync-fallback transfer", + ) + _run_async_transfer_items( + engine_wrapper=engine_wrapper, + by_session=by_session, + items=async_items, + session_debug_info=session_debug_info, + session_transfer_s=session_transfer_s, + bucket_idx=bucket_idx, + ) + else: + _run_sync_transfer_items( + engine_wrapper=engine_wrapper, + by_session=by_session, + items=all_large_items, + session_debug_info=session_debug_info, + session_transfer_s=session_transfer_s, + bucket_idx=bucket_idx, + label="batch_transfer_sync", + ) + + small_bytes, num_small_buffers = _transfer_small_entries( + engine_wrapper=engine_wrapper, + small_session_data=small_session_data, + session_debug_info=session_debug_info, + small_register_ptrs=small_register_ptrs, + small_register_lens=small_register_lens, + session_bytes=session_bytes, + session_transfer_s=session_transfer_s, + bucket_idx=bucket_idx, + ) + bucket_bytes += small_bytes + + timing.transfer_s = time.perf_counter() - t0 + timing.nbytes = bucket_bytes + timing.num_large_buffers = num_large_buffers + timing.num_small_buffers = num_small_buffers + timing.session_bytes = session_bytes + timing.session_transfer_s = session_transfer_s + if log_bucket_details: + logger.info( + "[P2P] bucket %d: %.1f MB, register=%.1f ms, transfer=%.1f ms, deregister=%.1f ms, throughput=%.1f MB/s", + bucket_idx, + timing.nbytes / 1e6, + timing.register_s * 1e3, + timing.transfer_s * 1e3, + timing.deregister_s * 1e3, + timing.throughput_mb_s, + ) + + +class P2PTransportBackend(WeightTransportBackend): + """RDMA P2P weight transport via the Mooncake TransferEngine. + + See module docstring for the architecture. The backend assumes: + + * The SGLang receiver exposes ``transport="p2p"`` on + ``/prepare_weights_update`` and returns ``tensor_map`` + + ``receiver_transfer_engine_infos``. + * The local training rank can construct a Mooncake TransferEngine + (mooncake-transfer-engine package installed, IB devices visible). + """ + + def __init__(self, config: TransportConfig, **kwargs: Any) -> None: + super().__init__(config) + # backend_config carries optional Mooncake engine setup overrides. + be_cfg = config.backend_config or {} + self._engine = None # MooncakeTransferEngine (lazy-imported) + # tensor_map[name] -> list of receiver locator dicts. + self._tensor_map: Dict[str, List[Dict[str, Any]]] = {} + # receiver_session_ids: list of session_id strings, one per inference TP rank. + self._receiver_session_ids: List[str] = [] + self._session_debug_info: Dict[str, Dict[str, Any]] = {} + # Warm-prepare coordination. When every trainer rank already has a + # cached tensor_map, rank 0 asks SGLang to prepare the update without + # returning the 100k+ locator JSON again, then broadcasts a tiny reuse + # marker instead of the full map. + self._prefer_cached_prepare: bool = False + self._last_prepare_returned_tensor_map: bool = False + # Source-side memory regions registered with Mooncake. The current + # async path registers CPU scratch pools once and registers small GPU + # entries per bucket in the worker. + self._registered_source_ptrs: List[int] = [] + self._registered_intervals: List[Tuple[int, int]] = [] + # Ring of CPU-pinned scratch pools for async transfer + # pipelining. While the worker thread runs Mooncake on pool A's + # bytes, the main thread stages layer N+1 into pool B (or C, D, + # etc.) — this hides per-bucket Mooncake latency behind the + # trainer-side FSDP unshard/extract work. Each pool is + # registered with Mooncake once at first use and reused for + # every subsequent bucket. + # + # Default 2 pools (ping-pong). Raising XORL_P2P_NUM_POOLS and + # XORL_P2P_MOONCAKE_WORKERS can hide per-call latency, but on + # the scaled TP2 layout 3-4 pools/workers regressed due to + # staging/NIC contention; treat higher values as experiments. + n_pools = max(1, int(os.environ.get("XORL_P2P_NUM_POOLS", "2"))) + self._n_pools = n_pools + self._cpu_scratch_pool_bytes: int = int(be_cfg.get("cpu_scratch_pool_bytes", _CPU_SCRATCH_POOL_BYTES)) + self._cpu_pool_min_bytes: int = int(be_cfg.get("cpu_pool_min_bytes", _CPU_POOL_MIN_BYTES)) + self._persist_small_registration: bool = _persist_small_registration_enabled() + self._cpu_scratch_pools: List[Optional[torch.Tensor]] = [None] * n_pools + self._cpu_scratch_pool_ptrs: List[int] = [0] * n_pools + self._cpu_scratch_pool_nbytes: int = 0 + self._cpu_pool_idx: int = 0 + self._cpu_pool_pending_futures: List[Optional[Future]] = [None] * n_pools + # Worker pool for async Mooncake calls. Default workers is 2; the + # number can be tuned with XORL_P2P_MOONCAKE_WORKERS. + self._transfer_executor: Optional[ThreadPoolExecutor] = None + # Per-bucket timings collected across this sync. Populated by + # transfer_bucket; read out by the caller (e.g. the e2e harness) + # for a wall-time breakdown vs. the NCCL backend. + self._bucket_timings: List[_BucketTiming] = [] + self._log_bucket_details: bool = _env_flag("XORL_P2P_LOG_BUCKET_DETAILS", False) + self._collect_transfer_debug: bool = _env_flag("XORL_P2P_TRANSFER_DEBUG", False) + self._transfer_chunk: int = _env_int("XORL_P2P_MOONCAKE_TRANSFER_CHUNK", 1) + # Stable group name passed back in /complete_weights_update. + self._group_name = config.group_name + self._hostname: Optional[str] = be_cfg.get("hostname") + self._resolved_hostname: Optional[str] = None + self._gpu_id: int = be_cfg.get("gpu_id", 0) + self._ib_device: Optional[str] = be_cfg.get("ib_device") + self._run_post_process_weights: bool = bool(be_cfg.get("run_post_process_weights", False)) + # ---- Direct EP / multi-sender configuration ---- + # When True, this backend is used in a multi-rank training setup + # where every training rank sends its own slice in parallel + # (the topology that produces the lmsys article's 7x speedup). + # The handler is responsible for invoking transfer_bucket on + # every rank in sender_ranks and only with that rank's params. + self._direct_ep_transfer: bool = bool(be_cfg.get("direct_ep_transfer", False)) + sender_ranks = be_cfg.get("sender_ranks") + self._explicit_sender_rank_order: Optional[Tuple[int, ...]] = None + self._explicit_sender_ranks: Optional[FrozenSet[int]] = None + if sender_ranks is not None: + self._explicit_sender_rank_order = tuple(dict.fromkeys(int(rank) for rank in sender_ranks)) + self._explicit_sender_ranks = frozenset(self._explicit_sender_rank_order) + self._process_group = be_cfg.get("process_group") + self._direct_ep_size: int = int(be_cfg.get("direct_ep_size", 0) or 0) + self._direct_ep_dense_sharding: bool = bool( + be_cfg.get("direct_ep_dense_sharding", _env_flag("XORL_P2P_DIRECT_EP_DENSE_SHARDING")) + ) + self._sender_ep_ranks: Dict[int, int] = {} + for item in be_cfg.get("sender_ep_ranks") or (): + try: + sender_rank, ep_rank = item + except (TypeError, ValueError): + logger.warning("[P2P] ignoring malformed sender_ep_ranks entry %r", item) + continue + self._sender_ep_ranks[int(sender_rank)] = int(ep_rank) + # Optional per-rank predicate. When set, transfer_bucket filters + # locator entries to only those that belong to *this* rank. + # Receives the locator dict; returns True if this rank should + # ship the corresponding slice. Default: ship everything (the + # single-sender / rank-0-only path). + self._rank_filter = be_cfg.get("rank_filter") + # The rank index to claim ownership for in multi-sender mode. + # Defaults to TransportConfig.training_rank. + self._rank_index: int = int(be_cfg.get("rank_index", config.training_rank)) + # World size of the trainer side. Only matters when + # direct_ep_transfer is on (every trainer rank ships its own + # slice in parallel). Fall back to TransportConfig.training_world_size + # — this is what the production handler sets — instead of a literal + # 1, otherwise sender_ranks silently degrades to {0} and the + # handler routes every non-rank-0 trainer through the gather/ + # broadcast fallback. + self._world_size: int = int(be_cfg.get("world_size", config.training_world_size or 1)) + self._last_prepare_tensor_map_endpoint_indices: set[int] = set() + + # ------------------------------------------------------------------ + # Lifecycle + # ------------------------------------------------------------------ + + @staticmethod + def _record_matches_endpoint(record: Dict[str, Any], endpoint_idx: int, num_endpoints: int) -> bool: + record_endpoint_idx = record.get("endpoint_idx") + if record_endpoint_idx is None: + return num_endpoints == 1 and endpoint_idx == 0 + try: + return int(record_endpoint_idx) == endpoint_idx + except (TypeError, ValueError): + return False + + @classmethod + def _drop_endpoint_records( + cls, + records: List[Dict[str, Any]], + endpoint_idx: int, + num_endpoints: int, + ) -> List[Dict[str, Any]]: + return [record for record in records if not cls._record_matches_endpoint(record, endpoint_idx, num_endpoints)] + + @classmethod + def _drop_endpoint_locators( + cls, + tensor_map: Dict[str, List[Dict[str, Any]]], + endpoint_idx: int, + num_endpoints: int, + ) -> Dict[str, List[Dict[str, Any]]]: + updated: Dict[str, List[Dict[str, Any]]] = {} + for name, locators in tensor_map.items(): + kept = cls._drop_endpoint_records(locators, endpoint_idx, num_endpoints) + if kept: + updated[name] = kept + return updated + + @classmethod + def _session_ids_for_endpoint( + cls, + records: List[Dict[str, Any]], + endpoint_idx: int, + num_endpoints: int, + ) -> set[str]: + return { + str(record["session_id"]) + for record in records + if record.get("session_id") and cls._record_matches_endpoint(record, endpoint_idx, num_endpoints) + } + + def _receiver_session_infos(self) -> List[Dict[str, Any]]: + infos: List[Dict[str, Any]] = [] + for sid in self._receiver_session_ids: + info = dict(self._session_debug_info.get(str(sid), {"session_id": sid})) + info.setdefault("session_id", sid) + infos.append(info) + return infos + + @staticmethod + def _expert_index_from_name(name: str) -> Optional[int]: + parts = name.split(".") + for idx, part in enumerate(parts[:-1]): + if part != "experts": + continue + try: + return int(parts[idx + 1]) + except ValueError: + return None + return None + + @classmethod + def _experts_per_ep( + cls, + tensor_map: Dict[str, List[Dict[str, Any]]], + ep_size: int, + ) -> Optional[int]: + if ep_size <= 1: + return None + expert_indices = { + expert_idx + for name in tensor_map + for expert_idx in (cls._expert_index_from_name(name),) + if expert_idx is not None + } + if not expert_indices: + return None + total_experts = max(expert_indices) + 1 + if expert_indices != set(range(total_experts)): + logger.warning("[P2P] expert tensor_map names are not contiguous; keeping full map on each sender") + return None + if total_experts % ep_size != 0: + logger.warning( + "[P2P] total_experts=%d is not divisible by direct_ep_size=%d; keeping full map on each sender", + total_experts, + ep_size, + ) + return None + return total_experts // ep_size + + @classmethod + def _endpoint_indices_for_tensor_map( + cls, + tensor_map: Dict[str, List[Dict[str, Any]]], + num_endpoints: int, + ) -> set[int]: + endpoint_indices: set[int] = set() + for locators in tensor_map.values(): + for loc in locators: + endpoint_idx = loc.get("endpoint_idx") + if endpoint_idx is None and num_endpoints == 1: + endpoint_indices.add(0) + continue + try: + endpoint_indices.add(int(endpoint_idx)) + except (TypeError, ValueError): + continue + return endpoint_indices + + @staticmethod + def _copy_locator_list_for_scatter( + locators: List[Dict[str, Any]], + copy_mode: str, + ) -> List[Dict[str, Any]]: + if copy_mode == "deep": + return [dict(loc) for loc in locators] + if copy_mode == "none": + return locators + # Default: keep an independent list per scatter payload without + # duplicating every immutable locator dict. Nonzero ranks copy the + # dicts again when adopting/merging prepared state. + return list(locators) + + @staticmethod + def _scatter_locator_copy_mode() -> str: + if "XORL_P2P_SCATTER_REUSE_LOCATORS" in os.environ: + if _env_flag("XORL_P2P_SCATTER_REUSE_LOCATORS", False): + return "none" + if "XORL_P2P_SCATTER_COPY_MODE" not in os.environ: + return "list" + raw_mode = os.environ.get("XORL_P2P_SCATTER_COPY_MODE") + if raw_mode is None: + # Fast path: locator lists/dicts are read-only after SGLang prepare, + # and scatter_object_list serializes each recipient payload anyway. + return "none" + raw = raw_mode.strip().lower() + if raw in {"deep", "dict", "dicts"}: + return "deep" + if raw in {"list", "shallow", "lists"}: + return "list" + if raw in {"none", "reuse"}: + return "none" + logger.warning("[P2P] invalid XORL_P2P_SCATTER_COPY_MODE=%r; using reuse", raw_mode) + return "none" + + def _filter_tensor_map_for_sender( + self, + tensor_map: Dict[str, List[Dict[str, Any]]], + sender_rank: int, + *, + experts_per_ep: Optional[int] = None, + locator_copy_mode: str = "list", + ) -> Dict[str, List[Dict[str, Any]]]: + sender_ep_rank = self._sender_ep_ranks.get(int(sender_rank)) + if experts_per_ep is None: + experts_per_ep = self._experts_per_ep(tensor_map, self._direct_ep_size) + if sender_ep_rank is None or experts_per_ep is None: + return { + name: self._copy_locator_list_for_scatter(locators, locator_copy_mode) + for name, locators in tensor_map.items() + } + + keep_dense = int(sender_rank) == 0 + filtered: Dict[str, List[Dict[str, Any]]] = {} + for name, locators in tensor_map.items(): + expert_idx = self._expert_index_from_name(name) + if expert_idx is None: + if self._direct_ep_dense_sharding: + if not self.should_send_dense_param(name, int(sender_rank)): + continue + elif not keep_dense: + continue + elif expert_idx // experts_per_ep != sender_ep_rank: + continue + filtered[name] = self._copy_locator_list_for_scatter(locators, locator_copy_mode) + return filtered + + def should_extract_dense_params_on_rank(self, rank: int) -> bool: + return self._direct_ep_dense_sharding and int(rank) in self.sender_ranks + + @staticmethod + def _dense_owner_key(name: str) -> str: + """Canonicalize equivalent trainer/SGLang dense names for sharding. + + The handler sees fused trainer names before ``_unfuse_for_inference`` + (for example ``qkv_proj`` and ``gate_up_proj``), while the receiver + tensor map sees the split SGLang names (``q_proj``/``k_proj``/``v_proj`` + and ``gate_proj``/``up_proj``). Hashing the canonical fused key keeps + extraction, post-unfuse filtering, and tensor-map filtering aligned. + """ + if P2PTransportBackend._expert_index_from_name(name) is not None: + return name + for split_name in (".q_proj.", ".k_proj.", ".v_proj."): + if split_name in name: + return name.replace(split_name, ".qkv_proj.", 1) + for split_name in (".gate_proj.", ".up_proj."): + if split_name in name: + return name.replace(split_name, ".gate_up_proj.", 1) + return name + + def should_send_dense_param(self, name: str, rank: int) -> bool: + """Return whether ``rank`` owns this dense parameter in direct-EP mode. + + MoE expert names are never treated as dense here. With dense sharding + disabled, rank 0 remains the sole dense sender, matching the historical + path. With sharding enabled, names are assigned deterministically across + the explicit sender rank order using a canonical fused dense name so + handler extraction and tensor-map filtering make the same decision on + every rank. + """ + rank = int(rank) + if self._expert_index_from_name(name) is not None: + return False + if not (self._direct_ep_transfer and self._world_size > 1 and self._direct_ep_dense_sharding): + return rank == 0 + sender_order = self.sender_rank_order + if not sender_order: + return rank == 0 + owner_key = self._dense_owner_key(name) + owner = sender_order[zlib.crc32(owner_key.encode("utf-8")) % len(sender_order)] + return rank == int(owner) + + def filter_dense_buffer_for_rank( + self, + buffer: List[Tuple[str, torch.Tensor]], + rank: int, + ) -> List[Tuple[str, torch.Tensor]]: + if not (self._direct_ep_transfer and self._world_size > 1): + return buffer if int(rank) == 0 else [] + if not self._direct_ep_dense_sharding: + return buffer if int(rank) == 0 else [] + return [(name, tensor) for name, tensor in buffer if self.should_send_dense_param(name, int(rank))] + + def _can_scatter_filtered_tensor_maps(self) -> bool: + return ( + self._direct_ep_transfer + and self._world_size > 1 + and self._explicit_sender_rank_order is not None + and bool(self._sender_ep_ranks) + and self._direct_ep_size > 1 + and hasattr(dist, "scatter_object_list") + ) + + def _initialize_payloads_for_sender_order(self) -> List[Any]: + session_infos = self._receiver_session_infos() + returned_endpoint_indices = set(self._last_prepare_tensor_map_endpoint_indices) + all_endpoint_indices = set(range(len(self.config.endpoints))) + if returned_endpoint_indices and returned_endpoint_indices != all_endpoint_indices: + kind = "merge_tensor_map" + else: + kind = "tensor_map_with_infos" + + experts_per_ep = self._experts_per_ep(self._tensor_map, self._direct_ep_size) + locator_copy_mode = self._scatter_locator_copy_mode() + payloads: List[Any] = [] + for sender_rank in self.sender_rank_order: + if int(sender_rank) == 0: + payloads.append(("rank0_ready",)) + continue + sender_tensor_map = self._filter_tensor_map_for_sender( + self._tensor_map, + sender_rank, + experts_per_ep=experts_per_ep, + locator_copy_mode=locator_copy_mode, + ) + if kind == "merge_tensor_map": + payloads.append( + ( + kind, + sender_tensor_map, + session_infos, + tuple(sorted(returned_endpoint_indices)), + ) + ) + else: + payloads.append((kind, sender_tensor_map, session_infos)) + return payloads + + def adopt_prepared_state( + self, + tensor_map: Dict[str, List[Dict[str, Any]]], + receiver_session_ids: List[str], + receiver_session_infos: Optional[List[Dict[str, Any]]] = None, + ) -> bool: + """Multi-sender hook: take the tensor_map + receiver session ids + that rank 0 obtained from ``/prepare_weights_update`` and stand up + the local Mooncake engine without doing the HTTP call again. + + Use this on every non-rank-0 sender. Rank 0 still calls + :meth:`initialize` to drive the prepare handshake; it then + broadcasts the tensor_map and session ids (e.g. via + ``torch.distributed.broadcast_object_list``) and the other + ranks call this. Each rank ends up with its own local engine + bound to its own GPU/HCA but a shared view of the receiver's + registered memory. + """ + # Reuse a cached engine if this backend is being re-initialized + # via the handler-side cache. + if self._engine is None: + self._engine = self._make_local_engine() + if self._engine is None: + return False + self._tensor_map = dict(tensor_map) + self._receiver_session_ids = list(receiver_session_ids) + if receiver_session_infos is not None: + self._session_debug_info = { + str(info["session_id"]): dict(info) for info in receiver_session_infos if info.get("session_id") + } + else: + self._session_debug_info = { + sid: {"session_id": sid, "adopted_from_rank0": True} for sid in self._receiver_session_ids + } + return True + + def merge_prepared_state( + self, + tensor_map: Dict[str, List[Dict[str, Any]]], + receiver_session_infos: List[Dict[str, Any]], + endpoint_indices: Tuple[int, ...], + ) -> bool: + if self._engine is None: + self._engine = self._make_local_engine() + if self._engine is None: + return False + + num_endpoints = len(self.config.endpoints) + updated = {name: [dict(loc) for loc in locators] for name, locators in self._tensor_map.items()} + indices = set(endpoint_indices) or self._endpoint_indices_for_tensor_map(tensor_map, num_endpoints) + for endpoint_idx in indices: + updated = self._drop_endpoint_locators(updated, endpoint_idx, num_endpoints) + for name, locators in tensor_map.items(): + updated.setdefault(name, []).extend(dict(loc) for loc in locators) + self._tensor_map = updated + self._receiver_session_ids = [ + str(info["session_id"]) for info in receiver_session_infos if info.get("session_id") + ] + self._session_debug_info = { + str(info["session_id"]): dict(info) for info in receiver_session_infos if info.get("session_id") + } + return True + + def initialize(self) -> bool: + # Multi-sender: rank 0 drives the HTTP prepare and broadcasts the + # result to every other sender. Non-zero ranks adopt that shared + # state and stand up their own local Mooncake engine. + if self._direct_ep_transfer and self._world_size > 1: + return self._initialize_multi_sender() + return self._initialize_single_sender() + + def _initialize_multi_sender(self) -> bool: + if not dist.is_available() or not dist.is_initialized(): + logger.error( + "[P2P] direct_ep_transfer set but torch.distributed is not " + "initialized. Initialize a process group first (the handler " + "does this for FSDP/EP)." + ) + return False + + is_rank0 = self._rank_index == 0 + group = self._process_group + group_world_size = len(self.sender_rank_order) + has_cached_state = bool(self._tensor_map and self._receiver_session_ids) + cached_states: List[bool] = [False] * group_world_size + try: + dist.all_gather_object(cached_states, has_cached_state, group=group) + except Exception as e: + logger.warning(f"[P2P] cached-state all_gather failed; using full prepare: {e}") + cached_states = [False] * group_world_size + self._prefer_cached_prepare = all(cached_states) + + if _env_flag("XORL_P2P_PREINIT_NONZERO_ENGINES") and not is_rank0 and self._engine is None: + logger.info("[P2P] pre-initializing local Mooncake engine while rank 0 prepares receivers") + self._engine = self._make_local_engine() + + payload: List[Any] = [None] + scatter_payloads: Optional[List[Any]] = None + use_scatter_payloads = self._can_scatter_filtered_tensor_maps() + if is_rank0: + ok = self._initialize_single_sender() + if ok: + if self._last_prepare_returned_tensor_map: + if use_scatter_payloads: + scatter_payloads = self._initialize_payloads_for_sender_order() + payload[0] = ("scatter_payloads",) + else: + payload[0] = ("tensor_map", self._tensor_map, list(self._receiver_session_ids)) + else: + payload[0] = ("reuse_cached", list(self._receiver_session_ids)) + else: + payload[0] = None + # All ranks synchronize on init payload delivery. When direct-EP + # sender mappings are available, scatter per-sender tensor maps instead + # of broadcasting the full 1M+ locator map to every sender. + if use_scatter_payloads: + scatter_input: Optional[List[Any]] = None + if is_rank0: + scatter_input = scatter_payloads if scatter_payloads is not None else [payload[0]] * group_world_size + if len(scatter_input) != group_world_size: + logger.error( + "[P2P] scatter payload count %d does not match sender group size %d", + len(scatter_input), + group_world_size, + ) + scatter_input = [None] * group_world_size + dist.scatter_object_list(payload, scatter_input, src=0, group=group) + else: + # payload[0] is None on failure so non-zero ranks can short-circuit + # cleanly. + dist.broadcast_object_list(payload, src=0, group=group) + local_ok = True + if payload[0] is None: + if not is_rank0: + logger.error("[P2P] rank 0 reported initialize() failure") + local_ok = False + elif not is_rank0: + kind = payload[0][0] + if kind == "reuse_cached": + if not has_cached_state or self._engine is None: + logger.error("[P2P] rank 0 requested cached prepare reuse, but this rank has no cached state") + local_ok = False + elif sids := payload[0][1]: + self._receiver_session_ids = list(sids) + self._session_debug_info = { + sid: {"session_id": sid, "adopted_from_rank0": True, "cached_prepare": True} + for sid in self._receiver_session_ids + } + elif kind == "tensor_map": + _, tmap, sids = payload[0] + if not self.adopt_prepared_state(tmap, sids): + local_ok = False + elif kind == "tensor_map_with_infos": + _, tmap, infos = payload[0] + sids = [str(info["session_id"]) for info in infos if info.get("session_id")] + if not self.adopt_prepared_state(tmap, sids, receiver_session_infos=infos): + local_ok = False + elif kind == "merge_tensor_map": + _, tmap, infos, endpoint_indices = payload[0] + if not self.merge_prepared_state(tmap, infos, tuple(int(idx) for idx in endpoint_indices)): + local_ok = False + else: + logger.error(f"[P2P] unknown initialize payload kind: {kind!r}") + local_ok = False + + init_results: List[bool] = [False] * group_world_size + try: + dist.all_gather_object(init_results, bool(local_ok), group=group) + except Exception as e: + logger.error(f"[P2P] direct-EP initialize result all_gather failed: {e}") + return False + if not all(init_results): + failed_ranks = [idx for idx, ok in enumerate(init_results) if not ok] + logger.error(f"[P2P] direct-EP initialize failed on ranks {failed_ranks}") + return False + return True + + def _initialize_single_sender(self) -> bool: + cfg = self.config + if not cfg.endpoints: + logger.error("[P2P] No endpoints provided in TransportConfig") + return False + + # If we're being reused via the handler-side cache, the + # engine is already constructed and the CPU scratch pools are + # already registered. Skip both — they're the slow steps. + if self._engine is None: + self._engine = self._make_local_engine() + if self._engine is None: + return False + + sender_session_id = self._engine.get_session_id() + sender_info = { + "session_id": sender_session_id, + "hostname": self._resolved_hostname or self._hostname, + "gpu_id": self._gpu_id, + "ib_device": self._engine.get_ib_device(), + "training_rank": cfg.training_rank, + } + + # Build prepare buckets once — we send them up front so SGLang + # can size its registration / state. + # For the single-sender variant we issue the prepare against each + # endpoint and merge the returned tensor_maps. Endpoints are + # independent, so fan out concurrently; on the 16-endpoint TP2 layout + # this keeps cached prepare from becoming a serialized HTTP/JSON tail. + has_cached_prepare_state = bool(self._tensor_map and self._receiver_session_ids) + if not (self._direct_ep_transfer and self._world_size > 1): + self._prefer_cached_prepare = has_cached_prepare_state + request_cached_prepare = self._prefer_cached_prepare and has_cached_prepare_state + self._last_prepare_returned_tensor_map = False + self._last_prepare_tensor_map_endpoint_indices = set() + num_endpoints = len(cfg.endpoints) + prepare_workers = min( + num_endpoints, + _env_int("XORL_P2P_PREPARE_WORKERS", min(32, num_endpoints)), + ) + + def _prepare_endpoint(ep_idx: int, ep: Any, cached_prepare: bool) -> Tuple[int, Any, Dict[str, Any]]: + url = f"http://{ep.host}:{ep.port}/prepare_weights_update" + payload = { + "buckets": [], # buckets are not used in the p2p path + "num_buckets": 0, + "group_name": cfg.group_name, + "transport": "p2p", + "sender_transfer_engine_info": sender_info, + } + if cached_prepare: + payload["p2p_return_tensor_map"] = False + try: + resp = requests.post(url, json=payload, timeout=_prepare_timeout_seconds()) + except requests.RequestException as e: + raise RuntimeError(f"/prepare_weights_update to {ep.host}:{ep.port} failed: {e}") from e + if cached_prepare and resp.status_code in (400, 422): + retry_payload = dict(payload) + retry_payload.pop("p2p_return_tensor_map", None) + logger.warning( + f"[P2P] cached prepare was rejected by {ep.host}:{ep.port}; " + "retrying this endpoint with full tensor_map response" + ) + try: + resp = requests.post(url, json=retry_payload, timeout=_prepare_timeout_seconds()) + except requests.RequestException as e: + raise RuntimeError(f"/prepare_weights_update retry to {ep.host}:{ep.port} failed: {e}") from e + if resp.status_code != 200: + raise RuntimeError( + f"/prepare_weights_update returned {resp.status_code} from {ep.host}:{ep.port}: {resp.text}" + ) + try: + body = resp.json() + except ValueError as e: + raise RuntimeError(f"/prepare_weights_update from {ep.host}:{ep.port} returned non-JSON") from e + if not body.get("success", False): + raise RuntimeError(f"prepare failed at {ep.host}:{ep.port}: {body.get('message')}") + return ep_idx, ep, body + + tensor_map_endpoint_indices: set[int] = set() + while True: + if request_cached_prepare: + # Cached prepare is the hot warm-sync path. Most warm prepares + # return no tensor_map from any endpoint, so avoid copying the + # large locator map unless an endpoint actually refreshes its + # locators below. + merged_tensor_map: Dict[str, List[Dict[str, Any]]] = self._tensor_map + else: + merged_tensor_map = {} + merged_receiver_infos: List[Dict[str, Any]] = [] + if request_cached_prepare: + for sid_idx, sid in enumerate(self._receiver_session_ids): + cached_info = dict(self._session_debug_info.get(str(sid), {"session_id": sid})) + cached_info.setdefault("session_id", sid) + if ( + "endpoint_idx" not in cached_info + and num_endpoints > 1 + and len(self._receiver_session_ids) == num_endpoints + ): + cached_info["endpoint_idx"] = sid_idx + merged_receiver_infos.append(cached_info) + + restart_full_prepare = False + tensor_map_endpoint_indices = set() + try: + if prepare_workers == 1: + endpoint_results = [ + _prepare_endpoint(ep_idx, ep, request_cached_prepare) for ep_idx, ep in enumerate(cfg.endpoints) + ] + else: + with ThreadPoolExecutor(max_workers=prepare_workers, thread_name_prefix="p2p-prepare") as executor: + futures = [ + executor.submit(_prepare_endpoint, ep_idx, ep, request_cached_prepare) + for ep_idx, ep in enumerate(cfg.endpoints) + ] + endpoint_results = [future.result() for future in as_completed(futures)] + except Exception as e: + logger.error(f"[P2P] prepare fanout failed: {e}") + return False + + for ep_idx, ep, body in sorted(endpoint_results, key=lambda item: item[0]): + ep_tensor_map = body.get("tensor_map") or {} + ep_receiver_infos = body.get("receiver_transfer_engine_infos") or [] + cached_sessions = self._session_ids_for_endpoint(merged_receiver_infos, ep_idx, num_endpoints) + returned_sessions = {str(info["session_id"]) for info in ep_receiver_infos if info.get("session_id")} + if ( + request_cached_prepare + and not ep_tensor_map + and cached_sessions + and returned_sessions + and returned_sessions != cached_sessions + ): + logger.warning( + f"[P2P] receiver sessions changed at {ep.host}:{ep.port}; " + "restarting prepare for all endpoints with full tensor_map response" + ) + request_cached_prepare = False + self._last_prepare_returned_tensor_map = False + restart_full_prepare = True + break + + if ep_tensor_map: + self._last_prepare_returned_tensor_map = True + tensor_map_endpoint_indices.add(ep_idx) + merged_tensor_map = self._drop_endpoint_locators(merged_tensor_map, ep_idx, num_endpoints) + for name, locators in ep_tensor_map.items(): + # Tag each locator with its source endpoint so transfer_bucket + # knows where to reach it. + for loc in locators: + loc = dict(loc) + loc["endpoint_idx"] = ep_idx + merged_tensor_map.setdefault(name, []).append(loc) + + if ep_receiver_infos: + merged_receiver_infos = self._drop_endpoint_records(merged_receiver_infos, ep_idx, num_endpoints) + for info in ep_receiver_infos: + info = dict(info) + info["endpoint_idx"] = ep_idx + merged_receiver_infos.append(info) + + if restart_full_prepare: + continue + break + self._last_prepare_tensor_map_endpoint_indices = set(tensor_map_endpoint_indices) + + if merged_tensor_map: + self._tensor_map = merged_tensor_map + elif not self._tensor_map: + logger.error("[P2P] prepare returned no tensor_map and no cached map is available") + return False + + if merged_receiver_infos: + self._receiver_session_ids = [ + str(info["session_id"]) for info in merged_receiver_infos if info.get("session_id") + ] + self._session_debug_info = { + str(info["session_id"]): dict(info) for info in merged_receiver_infos if info.get("session_id") + } + elif not self._receiver_session_ids: + logger.error("[P2P] prepare returned no receiver session ids and no cached sessions are available") + return False + + total_locators = sum(len(v) for v in self._tensor_map.values()) + logger.info( + f"[P2P] prepare ok: {len(self._tensor_map)} hf_names, " + f"{total_locators} locators across " + f"{len(self._receiver_session_ids)} receivers " + f"(cached_prepare={request_cached_prepare and not self._last_prepare_returned_tensor_map}, " + f"prepare_workers={prepare_workers}, " + f"tensor_map_endpoints={len(tensor_map_endpoint_indices)}/{num_endpoints})" + ) + return True + + def _ensure_cpu_scratch_pool(self) -> None: + """Lazy-init the CPU pinned scratch pools. + + Pools are allocated and registered with Mooncake only when a bucket + has at least one large entry that uses the CPU-source path. The + transfer executor is initialized separately so small-only buckets do + not allocate gigabytes of pinned memory. + """ + if self._cpu_scratch_pools[0] is not None: + return + if self._engine is None: + raise RuntimeError("[P2P] _ensure_cpu_scratch_pool called before initialize()") + n = self._cpu_scratch_pool_bytes + for i in range(self._n_pools): + # uint8 makes byte-level offset math straightforward. + pool = torch.empty(n, dtype=torch.uint8, pin_memory=True) + ptr = int(pool.data_ptr()) + t0 = time.perf_counter() + # Use the singular API for CPU pinned memory; it routes through a + # different Mooncake code path than the GPU-oriented batch register. + ret = self._engine.engine.register_memory(ptr, n) + dt = time.perf_counter() - t0 + if ret != 0: + raise RuntimeError(f"[P2P] CPU scratch pool {i} register_memory failed: ret={ret} ({n / 1e9:.2f} GB)") + self._cpu_scratch_pools[i] = pool + self._cpu_scratch_pool_ptrs[i] = ptr + self._registered_source_ptrs.append(ptr) + self._registered_intervals.append((ptr, ptr + n)) + logger.info( + f"[P2P] CPU scratch pool {i}/{self._n_pools}: registered {n / 1e9:.2f} GB " + f"pinned at 0x{ptr:x} in {dt * 1000:.1f} ms" + ) + self._cpu_scratch_pool_nbytes = n + self._ensure_transfer_executor() + + def _ensure_transfer_executor(self) -> None: + """Lazy-init the worker pool that runs Mooncake transfer calls.""" + if self._transfer_executor is not None: + return + # Number of concurrent Mooncake calls per rank. More than one worker + # lets slower per-call paths hide Mooncake latency behind subsequent + # staged buckets, especially when rank 0 has many small dense buckets. + # Set XORL_P2P_MOONCAKE_WORKERS to override. + n_workers = max(1, int(os.environ.get("XORL_P2P_MOONCAKE_WORKERS", "2"))) + self._transfer_executor = ThreadPoolExecutor(max_workers=n_workers, thread_name_prefix="p2p-mooncake") + if n_workers > 1: + logger.info(f"[P2P] using {n_workers} concurrent Mooncake workers") + + def _wait_all_pending(self) -> None: + """Block until every outstanding async transfer completes. + + Must be called before reading ``stats_summary`` or tearing down + the engine — timings and bucket_bytes are only valid after the + worker has populated them. + """ + first_error: Optional[Exception] = None + for i in range(len(self._cpu_pool_pending_futures)): + fut = self._cpu_pool_pending_futures[i] + if fut is None: + continue + try: + fut.result() + except Exception as e: + logger.error(f"[P2P] async transfer on pool {i} raised: {e}") + if first_error is None: + first_error = e + finally: + self._cpu_pool_pending_futures[i] = None + if first_error is not None: + raise first_error + + def flush_pending_transfers(self) -> None: + """Drain async transfers before the handler resumes inference + or measures wall time. Otherwise generation can read + partially-updated weights, since RDMA writes land directly in + ``param.data`` without going through ``/complete_weights_update``. + """ + self._wait_all_pending() + + def complete_sync(self) -> None: + """Per-sync teardown: drain in-flight transfers + send the + receiver-side completion RPC. Leaves engine + CPU scratch pools + + executor intact so the next sync can reuse them through the handler + cache. + + Safe to call multiple times after a successful transfer drain. Not safe + to call for a failed/partial sync: pending transfer errors are raised + before the receiver-side completion RPC is sent. + """ + cfg = self.config + # Drain async transfers before any teardown step; required even + # in cache-reuse mode because pending futures hold references to + # CPU pool slots that the next sync will overwrite. + self._wait_all_pending() + + # In multi-sender mode (direct_ep_transfer + world_size>1), only + # rank 0 drives the HTTP /complete_weights_update — the receiver + # services exactly one complete per sync, so non-zero ranks must + # skip it or they trip the "no update in progress" error. + skip_complete = self._direct_ep_transfer and self._world_size > 1 and self._rank_index != 0 + if not skip_complete: + be_cfg = cfg.backend_config or {} + flush_cache = bool(be_cfg.get("flush_cache", False)) + weight_version = be_cfg.get("weight_version") + tied_weight_aliases = be_cfg.get("p2p_tied_weight_aliases") or {} + complete_errors = [] + for ep in cfg.endpoints: + url = f"http://{ep.host}:{ep.port}/complete_weights_update" + payload = { + "group_name": cfg.group_name, + "flush_cache": flush_cache, + "transport": "p2p", + "run_post_process_weights": self._run_post_process_weights, + } + if weight_version is not None: + payload["weight_version"] = weight_version + if tied_weight_aliases: + payload["p2p_tied_weight_aliases"] = tied_weight_aliases + try: + resp = requests.post(url, json=payload, timeout=_HTTP_TIMEOUT_SECONDS) + if resp.status_code != 200: + complete_errors.append(f"{ep.host}:{ep.port} returned HTTP {resp.status_code}: {resp.text}") + continue + try: + body = resp.json() + except ValueError: + body = {} + if body and body.get("success") is False: + complete_errors.append( + f"{ep.host}:{ep.port} returned success=false: {body.get('message', body)}" + ) + except requests.RequestException as e: + complete_errors.append(f"{ep.host}:{ep.port} request failed: {e}") + if complete_errors: + raise RuntimeError("[P2P] /complete_weights_update failed: " + "; ".join(complete_errors)) + + # Keep the receiver tensor_map/session ids with the cached backend. + # The next sync can ask a warm SGLang receiver to skip returning the + # huge locator JSON and can reuse this local map instead. If the + # receiver restarted or rejected the cached prepare path, initialize() + # replaces this state from the full prepare response. + self._bucket_timings = [] + # CPU pool ping-pong cursor: reset so the next sync starts on + # pool 0 (deterministic, makes timing logs easier to read). + self._cpu_pool_idx = 0 + + @property + def is_alive(self) -> bool: + """True if engine + scratch pools are still allocated and + registered. Caller can use this to decide whether to reuse this + backend on the next sync.""" + return self._engine is not None + + def destroy(self, *, complete_receiver: bool = True) -> None: + """Full teardown: per-sync complete + release engine, executor, + and CPU scratch pools. After this the backend cannot be reused.""" + complete_error: Optional[Exception] = None + if complete_receiver: + try: + self.complete_sync() + except Exception as e: + complete_error = e + else: + try: + self._wait_all_pending() + except Exception as e: + logger.warning(f"[P2P] skipping receiver completion after failed/aborted sync: {e}") + self._bucket_timings = [] + self._cpu_pool_idx = 0 + + if self._transfer_executor is not None: + self._transfer_executor.shutdown(wait=True) + self._transfer_executor = None + + # Best-effort source-side deregistration. Includes the CPU + # pinned scratch pools (registered once at first transfer_bucket + # call) and any leftover per-bucket registrations from older + # code paths. + if self._engine is not None and self._registered_source_ptrs: + try: + self._engine.batch_deregister(self._registered_source_ptrs) + except Exception as e: + logger.warning(f"[P2P] batch_deregister of source pointers failed: {e}") + self._registered_source_ptrs = [] + self._registered_intervals = [] + # Release CPU pinned pools. PyTorch frees them on GC; we just + # drop the handles here for clarity. + self._cpu_scratch_pools = [None] * self._n_pools + self._cpu_scratch_pool_ptrs = [0] * self._n_pools + self._cpu_scratch_pool_nbytes = 0 + self._cpu_pool_idx = 0 + self._cpu_pool_pending_futures = [None] * self._n_pools + + self._engine = None + if complete_error is not None: + raise complete_error + + # ------------------------------------------------------------------ + # Transfer + # ------------------------------------------------------------------ + + def transfer_bucket( + self, + bucket: List[Tuple[str, torch.Tensor]], + *, + src_rank: int = 0, + flush_cache: bool = False, + weight_version: Optional[str] = None, + ) -> None: + if not self._direct_ep_transfer and src_rank != 0: + raise ValueError( + f"P2PTransportBackend rejects src_rank={src_rank} when " + "direct_ep_transfer is off (set " + "TransportConfig.backend_config['direct_ep_transfer']=True " + "for multi-sender mode)." + ) + if self._engine is None: + raise RuntimeError("P2P backend not initialized — call initialize() first") + + # weight_version + flush_cache get applied at /complete_weights_update. + # The handler signals these on the final bucket; stash them even if + # this particular rank has no locator-owned transfers. + if weight_version is not None: + self.config.backend_config["weight_version"] = weight_version + self.config.backend_config["flush_cache"] = bool( + flush_cache or self.config.backend_config.get("flush_cache", False) + ) + + # Two-pass design: + # + # Pass 1: collect (session_id, peer_ptr, nbytes, src_view) for + # every locator in the bucket. + # Pass 2 per session: sort by peer_ptr, stage src_views into the + # CPU pool in peer-ptr order so pool offsets match peer-ptr + # adjacency, then coalesce neighboring entries whose pool & peer + # ptrs are adjacent into single (src_ptr, peer_ptr, nbytes) tuples. + # Ship as one batch_transfer_sync per session. + # + # Why this matters: receiver locators emit per-expert HF names + # (e.g. experts.0.gate_proj.weight) that physically live in + # contiguous slices of receiver-internal tensors (w13_weight). + # Without coalescing, our trainer ships 96 buffers per layer's + # MoE batch, which exceeds Mooncake's per-call stability cap + # (~96-buffer ret=-1 mode observed empirically). After + # coalescing, a layer's 96 contiguous expert slices collapse + # to ~2 entries (w13 region, w2 region) per receiver-rank. + # This yields one entry per receiver-internal name, fully populated in + # CPU pinned source. + t_prepare = time.perf_counter() + pending: Dict[str, List[_PendingTransfer]] = {} + skipped_errors: List[str] = [] + for name, tensor in bucket: + locators = self._locators_for_source_name(name) + if not locators: + if self._should_skip_missing_tied_locator(name): + logger.info("[P2P] skipping missing tied receiver locator for %s", name) + continue + skipped_errors.append(f"{name!r}: no receiver locator") + continue + locators_for_rank = 0 + for loc in locators: + if self._rank_filter is not None and not self._rank_filter(loc): + continue + locators_for_rank += 1 + src_view = self._slice_source_for_locator(name, tensor, loc) + if src_view is None: + skipped_errors.append(f"{name!r}: receiver locator is incompatible with source tensor") + continue + src_nbytes = src_view.numel() * src_view.element_size() + expected = int(loc.get("nbytes", src_nbytes)) + if src_nbytes != expected: + skipped_errors.append( + f"[P2P] size mismatch for {name!r} after slicing: " + f"source={src_nbytes}, receiver expects={expected}. " + f"Check tensor map slice metadata." + ) + continue + session_id = loc.get("session_id") + if not session_id: + skipped_errors.append(f"{name!r}: receiver locator is missing session_id") + continue + pending.setdefault(session_id, []).append( + _PendingTransfer( + peer_ptr=int(loc["ptr"]), + nbytes=src_nbytes, + src_view=src_view, + name=name, + loc=loc, + ) + ) + if locators_for_rank == 0: + logger.debug("[P2P] no receiver locators owned by this sender for parameter %r", name) + + if skipped_errors: + preview = "; ".join(skipped_errors[:5]) + if len(skipped_errors) > 5: + preview += f"; ... {len(skipped_errors) - 5} more" + raise RuntimeError(f"[P2P] receiver tensor_map is incomplete or incompatible: {preview}") + + if not pending: + logger.debug("[P2P] transfer_bucket produced no transfers for this sender") + return + + timing = _BucketTiming() + self._bucket_timings.append(timing) + bucket_idx = len(self._bucket_timings) + timing.prepare_s = time.perf_counter() - t_prepare + + # Split entries into "large" (CPU pool path) and "small" + # (GPU-direct path with per-bucket register). Mooncake's CPU-source + # path on raw tiny buffers is unstable on our cluster, so only CUDA + # tensors below the threshold use GPU-direct. CPU tensors, including + # trainer-side FP8 scale tensors, are always staged through the + # pre-registered CPU pool. + pending_small: Dict[str, List[_PendingTransfer]] = {} + for sid, entries in list(pending.items()): + small: List[_PendingTransfer] = [] + large: List[_PendingTransfer] = [] + for e in entries: + if e.src_view.is_cuda and e.nbytes < self._cpu_pool_min_bytes: + small.append(e) + else: + large.append(e) + if small: + pending_small[sid] = small + if large: + pending[sid] = large + else: + pending.pop(sid, None) + + # First-call lazy init. Large entries need the registered CPU pool; + # small-only buckets only need the transfer executor. + t_pool_init = time.perf_counter() + if pending: + self._ensure_cpu_scratch_pool() + else: + self._ensure_transfer_executor() + timing.pool_init_s = time.perf_counter() - t_pool_init + + # Pick the current pool/future slot. If a prior transfer is still + # using it, block until it drains before overwriting the slot. + pool_idx = self._cpu_pool_idx + prior_future = self._cpu_pool_pending_futures[pool_idx] + t_pool_wait = time.perf_counter() + if prior_future is not None: + prior_future.result() # surface worker exceptions + self._cpu_pool_pending_futures[pool_idx] = None + timing.pool_wait_s = time.perf_counter() - t_pool_wait + pool = self._cpu_scratch_pools[pool_idx] if pending else None + pool_ptr = self._cpu_scratch_pool_ptrs[pool_idx] if pending else 0 + + # Hold references to staged sub-tensors and src views so the + # caller can reshard immediately after this method returns + # without freeing GPU memory the NIC is still reading. Both + # lists travel into the future closure and stay alive until + # the worker finishes. + slice_holds: List[torch.Tensor] = [] + src_view_holds: List[torch.Tensor] = [] + # Per-session transfer lists after coalescing. The debug list tracks + # which original locators contributed to each emitted Mooncake buffer. + by_session: Dict[str, Tuple[List[int], List[int], List[int], Optional[List[_TransferDebugSample]]]] = {} + scratch_offset_bytes = 0 + staged_sources: Dict[Tuple[int, Tuple[int, ...], Tuple[int, ...], str, int], int] = {} + unique_staged_bytes = 0 + reused_staged_bytes = 0 + total_pre_coalesce = 0 + total_post_coalesce = 0 + t_stage = time.perf_counter() + + for session_id, entries in pending.items(): + # Sort by peer_ptr so adjacent receiver-side regions cluster + # together for the coalesce pass. Source pool offsets are + # also assigned in this order, so pool adjacency matches + # peer adjacency exactly. + entries.sort(key=lambda e: e.peer_ptr) + total_pre_coalesce += len(entries) + + # Stage all src_views into the CPU pool in sorted order. + staged: List[_StagedTransfer] = [] + for e in entries: + if pool is None: + raise RuntimeError("[P2P] CPU scratch pool was not initialized") + source_key = _source_view_key(e.src_view, e.nbytes) + src_ptr = staged_sources.get(source_key) + if src_ptr is None: + element_size = max(1, int(e.src_view.element_size())) + scratch_offset_bytes = _align_up(pool_ptr + scratch_offset_bytes, element_size) - pool_ptr + if scratch_offset_bytes + e.nbytes > self._cpu_scratch_pool_nbytes: + # The scratch pool must hold the largest staged bucket for + # this sender. The default P2P MoE bucket cap is 2 GiB and + # the default pool is 4 GiB, so raising the bucket cap may + # require raising XORL_P2P_CPU_SCRATCH_POOL_BYTES too. + raise RuntimeError( + f"[P2P] CPU scratch pool exhausted: bucket needs " + f">{scratch_offset_bytes + e.nbytes} bytes but pool " + f"is {self._cpu_scratch_pool_nbytes} bytes. Increase " + f"XORL_P2P_CPU_SCRATCH_POOL_BYTES." + ) + slot_uint8 = pool[scratch_offset_bytes : scratch_offset_bytes + e.nbytes] + slot_view = slot_uint8.view(e.src_view.dtype).view(e.src_view.shape) + slot_view.copy_(e.src_view, non_blocking=True) + slice_holds.append(slot_view) + src_view_holds.append(e.src_view) + src_ptr = pool_ptr + scratch_offset_bytes + staged_sources[source_key] = src_ptr + unique_staged_bytes += e.nbytes + scratch_offset_bytes += e.nbytes + else: + reused_staged_bytes += e.nbytes + if e.src_view.is_cuda: + src_view_holds.append(e.src_view) + staged.append( + _StagedTransfer( + src_ptr=src_ptr, + peer_ptr=e.peer_ptr, + nbytes=e.nbytes, + memory_handle=_locator_memory_handle(e.loc), + name=e.name, + loc=e.loc, + ) + ) + + # Coalesce: walk staged in sorted-by-peer order, merging + # adjacent entries whose source pool ptr AND peer ptr are + # contiguous within the same receiver registration. After this + # pass, a layer's per-expert slices collapse to a small number + # of receiver-backed regions per session without asking + # Mooncake to write across registration boundaries. + src_ptrs: List[int] = [] + peer_ptrs: List[int] = [] + lens: List[int] = [] + debug_entries: Optional[List[_TransferDebugSample]] = [] if self._collect_transfer_debug else None + memory_handles: List[Optional[int]] = [] + for staged_entry in staged: + if ( + src_ptrs + and src_ptrs[-1] + lens[-1] == staged_entry.src_ptr + and peer_ptrs[-1] + lens[-1] == staged_entry.peer_ptr + and staged_entry.memory_handle is not None + and memory_handles[-1] == staged_entry.memory_handle + ): + lens[-1] += staged_entry.nbytes + if debug_entries is not None: + debug_entries[-1].add(staged_entry.name, staged_entry.loc, staged_entry.nbytes) + else: + src_ptrs.append(staged_entry.src_ptr) + peer_ptrs.append(staged_entry.peer_ptr) + lens.append(staged_entry.nbytes) + memory_handles.append(staged_entry.memory_handle) + if debug_entries is not None: + debug_sample = _TransferDebugSample() + debug_sample.add(staged_entry.name, staged_entry.loc, staged_entry.nbytes) + debug_entries.append(debug_sample) + total_post_coalesce += len(src_ptrs) + by_session[session_id] = (src_ptrs, peer_ptrs, lens, debug_entries) + + if self._log_bucket_details and total_post_coalesce < total_pre_coalesce: + logger.info( + f"[P2P] coalesced {total_pre_coalesce} entries → " + f"{total_post_coalesce} ({total_pre_coalesce / total_post_coalesce:.1f}x reduction)" + ) + if self._log_bucket_details and reused_staged_bytes: + logger.info( + "[P2P] staged-source reuse saved %.1f MB (unique_staged=%.1f MB, transfer_fanout=%.1f MB)", + reused_staged_bytes / 1e6, + unique_staged_bytes / 1e6, + (unique_staged_bytes + reused_staged_bytes) / 1e6, + ) + + # Build small-entries metadata on the main thread — Mooncake's + # batch_register doesn't enqueue CUDA work but + # _intervals_per_cuda_segment calls torch.cuda.memory_snapshot() + # which we keep on main as a precaution. The actual transfer + # work happens in the worker. + small_session_data: Dict[str, List[Tuple[int, int, int, Optional[_TransferDebugEntry]]]] = {} + small_register_ptrs: List[int] = [] + small_register_lens: List[int] = [] + if pending_small: + small_persistent_intervals: List[Tuple[int, int]] = [] + small_transient_intervals: List[Tuple[int, int]] = [] + for session_id, entries in pending_small.items(): + triples: List[Tuple[int, int, int, Optional[_TransferDebugEntry]]] = [] + for e in entries: + sv = e.src_view.contiguous() + src_view_holds.append(sv) + storage = sv.untyped_storage() + s_start = int(storage.data_ptr()) + s_end = s_start + int(storage.nbytes()) + # Persistently register only no-copy contiguous views backed + # by stable model-parameter storage. Temporary contiguous + # copies must stay per-bucket because they are released after + # the async worker drains. + can_persist = ( + self._persist_small_registration + and int(sv.data_ptr()) == int(e.src_view.data_ptr()) + and int(sv.untyped_storage().data_ptr()) == int(e.src_view.untyped_storage().data_ptr()) + ) + if can_persist: + small_persistent_intervals.append((s_start, s_end)) + else: + small_transient_intervals.append((s_start, s_end)) + debug_entry = ( + _transfer_debug_entry(e.name, e.loc, e.nbytes) if self._collect_transfer_debug else None + ) + triples.append((int(sv.data_ptr()), e.peer_ptr, e.nbytes, debug_entry)) + small_session_data[session_id] = triples + if small_persistent_intervals: + self._register_persistent_source_intervals(small_persistent_intervals, bucket_idx=bucket_idx) + if small_transient_intervals: + small_segs = self._intervals_per_cuda_segment(small_transient_intervals) + small_register_ptrs = [iv[0] for iv in small_segs] + small_register_lens = [iv[1] - iv[0] for iv in small_segs] + timing.stage_s = time.perf_counter() - t_stage + + # The CPU pool is permanently registered; small entries + # register/dereg in the worker. Both register_s and + # deregister_s stay zero on the main-thread accounting because + # all that work moves into transfer_s on the worker. + timing.register_s = 0.0 + timing.deregister_s = 0.0 + + # CUDA event records "all GPU work enqueued so far has + # completed". The worker waits on this before issuing Mooncake + # calls because Mooncake's NIC reads bypass CUDA streams entirely. + # CPU-only protocol tests use a no-op event. + if torch.cuda.is_available(): + copy_done_event = torch.cuda.Event() + copy_done_event.record() + else: + copy_done_event = _CompletedCudaEvent() + + if self._transfer_executor is None: + raise RuntimeError("[P2P] transfer executor was not initialized") + t_submit = time.perf_counter() + future = self._transfer_executor.submit( + _do_async_transfer, + engine_wrapper=self._engine, + copy_done_event=copy_done_event, + by_session=by_session, + small_session_data=small_session_data, + session_debug_info=self._session_debug_info, + small_register_ptrs=small_register_ptrs, + small_register_lens=small_register_lens, + chunk=self._transfer_chunk, + use_async_api=_async_api_enabled(cached_prepare=self._prefer_cached_prepare), + timing=timing, + bucket_idx=bucket_idx, + slice_holds=slice_holds, + src_view_holds=src_view_holds, + log_bucket_details=self._log_bucket_details, + ) + timing.submit_s = time.perf_counter() - t_submit + self._cpu_pool_pending_futures[pool_idx] = future + + # Round-robin to the next pool. With N pools and N workers, + # the main thread can stage up to N buckets ahead while workers + # drain them in parallel. + self._cpu_pool_idx = (pool_idx + 1) % self._n_pools + + # ------------------------------------------------------------------ + # Topology hints + # ------------------------------------------------------------------ + + @property + def bucket_timings(self) -> List[_BucketTiming]: + """Per-bucket wall-time breakdown collected during this sync.""" + # Async transfers may still be in flight; their timing fields + # are filled in by the worker thread. Drain before snapshotting + # so the caller sees a self-consistent list. + self._wait_all_pending() + return list(self._bucket_timings) + + def stats_summary(self) -> Dict[str, Any]: + """Aggregate timing summary for the most recent sync. + + Returns a dict with: ``num_buckets``, ``total_bytes``, ``total_s``, + main-thread staging fields, ``register_s``, ``transfer_s``, ``deregister_s``, + ``transfer_throughput_mb_s`` (transfer-only — excludes + main-thread/register/deregister), ``effective_throughput_mb_s`` + (bucket wall-time, the number to compare against NCCL). + """ + self._wait_all_pending() + timings = self._bucket_timings + bucket_transfer_s = [t.transfer_s for t in timings] + bucket_total_s = [t.total_s for t in timings] + total_bytes = sum(t.nbytes for t in timings) + prepare_s = sum(t.prepare_s for t in timings) + pool_init_s = sum(t.pool_init_s for t in timings) + pool_wait_s = sum(t.pool_wait_s for t in timings) + stage_s = sum(t.stage_s for t in timings) + submit_s = sum(t.submit_s for t in timings) + register_s = sum(t.register_s for t in timings) + transfer_s = sum(t.transfer_s for t in timings) + deregister_s = sum(t.deregister_s for t in timings) + main_thread_s = prepare_s + pool_init_s + pool_wait_s + stage_s + submit_s + total_s = sum(t.total_s for t in timings) + session_bytes: Dict[str, int] = {} + session_transfer_s: Dict[str, float] = {} + for timing in timings: + for session_id, nbytes in timing.session_bytes.items(): + session_bytes[session_id] = session_bytes.get(session_id, 0) + nbytes + for session_id, seconds in timing.session_transfer_s.items(): + session_transfer_s[session_id] = session_transfer_s.get(session_id, 0.0) + seconds + + def _percentile(values: List[float], percentile: float) -> float: + if not values: + return 0.0 + ordered = sorted(values) + index = min(len(ordered) - 1, max(0, int((percentile / 100.0) * len(ordered)))) + return ordered[index] + + top_sessions_by_transfer_s = [] + for session_id, seconds in session_transfer_s.items(): + nbytes = session_bytes.get(session_id, 0) + top_sessions_by_transfer_s.append( + { + "session_id": session_id, + "transfer_s": seconds, + "total_bytes": nbytes, + "throughput_mb_s": ((nbytes / 1e6) / seconds if seconds > 0 else 0.0), + } + ) + top_sessions_by_transfer_s.sort(key=lambda row: row["transfer_s"], reverse=True) + + slowest_buckets = [ + { + "bucket": bucket_idx, + "total_s": timing.total_s, + "main_thread_s": timing.main_thread_s, + "prepare_s": timing.prepare_s, + "pool_init_s": timing.pool_init_s, + "pool_wait_s": timing.pool_wait_s, + "stage_s": timing.stage_s, + "submit_s": timing.submit_s, + "transfer_s": timing.transfer_s, + "total_bytes": timing.nbytes, + "large_buffers": timing.num_large_buffers, + "small_buffers": timing.num_small_buffers, + } + for bucket_idx, timing in sorted( + enumerate(timings, start=1), + key=lambda item: item[1].total_s, + reverse=True, + )[:5] + ] + + return { + "num_buckets": float(len(timings)), + "total_bytes": float(total_bytes), + "prepare_s": prepare_s, + "pool_init_s": pool_init_s, + "pool_wait_s": pool_wait_s, + "stage_s": stage_s, + "submit_s": submit_s, + "main_thread_s": main_thread_s, + "register_s": register_s, + "transfer_s": transfer_s, + "deregister_s": deregister_s, + "total_s": total_s, + "transfer_throughput_mb_s": ((total_bytes / 1e6) / transfer_s if transfer_s > 0 else 0.0), + "effective_throughput_mb_s": ((total_bytes / 1e6) / total_s if total_s > 0 else 0.0), + "max_bucket_transfer_s": max(bucket_transfer_s) if bucket_transfer_s else 0.0, + "p50_bucket_transfer_s": _percentile(bucket_transfer_s, 50), + "p95_bucket_transfer_s": _percentile(bucket_transfer_s, 95), + "max_bucket_total_s": max(bucket_total_s) if bucket_total_s else 0.0, + "p50_bucket_total_s": _percentile(bucket_total_s, 50), + "p95_bucket_total_s": _percentile(bucket_total_s, 95), + "num_large_buffers": float(sum(t.num_large_buffers for t in timings)), + "num_small_buffers": float(sum(t.num_small_buffers for t in timings)), + "slowest_buckets": slowest_buckets, + "top_sessions_by_transfer_s": top_sessions_by_transfer_s[:5], + } + + @property + def sender_ranks(self) -> FrozenSet[int]: + if self._direct_ep_transfer and self._world_size > 1: + if self._explicit_sender_ranks is not None: + return self._explicit_sender_ranks + return frozenset(range(self._world_size)) + return frozenset({0}) + + @property + def sender_rank_order(self) -> Tuple[int, ...]: + if self._direct_ep_transfer and self._world_size > 1: + if self._explicit_sender_rank_order is not None: + return self._explicit_sender_rank_order + return tuple(range(self._world_size)) + return (0,) + + @property + def has_explicit_sender_ranks(self) -> bool: + return self._explicit_sender_ranks is not None + + @property + def supports_direct_ep_transfer(self) -> bool: + return self._direct_ep_transfer + + @property + def supports_direct_pp_transfer(self) -> bool: + # PP stage leaders still route through rank 0 in the handler. Keep + # this false until the PP-direct handler path is implemented. + return False + + # ------------------------------------------------------------------ + # Internals + # ------------------------------------------------------------------ + + def _intervals_per_cuda_segment( + self, + intervals: List[Tuple[int, int]], + ) -> List[Tuple[int, int]]: + """Constrain each interval to the CUDA-allocator's *active_allocated* + blocks. + + ``ibv_reg_mr`` (and Mooncake's wrapper) returns EFAULT when asked + to register a virtual-address range that crosses two distinct + physical mappings, or that includes pages currently cached/free + in PyTorch's caching allocator (those virtual addresses are reserved but not + backed by physical memory the IB driver can pin). + + Mirrors SGLang's ``register_memory_region_v2`` — walks the + ``memory_snapshot`` per-segment and within each segment accumulates + runs of contiguous ``active_allocated`` blocks that overlap the + candidate intervals, emitting one merged range per run. + """ + if not intervals: + return [] + try: + snapshot = torch.cuda.memory.memory_snapshot() + except Exception as e: + logger.warning( + f"[P2P] torch.cuda.memory.memory_snapshot failed: {e}; " + "falling back to raw intervals (may EFAULT on registration)." + ) + return intervals + + # Sort the candidate intervals once for an O(N+M) sweep. + sorted_candidates = sorted(intervals, key=lambda iv: iv[0]) + + def _overlaps_any(start: int, end: int) -> bool: + for cs, ce in sorted_candidates: + if ce <= start: + continue + if cs >= end: + return False + return True + return False + + # Register each active_allocated block separately. We previously + # tried merging adjacent blocks within a segment, but Mooncake's + # ibv_reg_mr returned EFAULT on the merged ranges (likely because + # the merged span crosses an internal allocator boundary that + # isn't representable as one MR). One-block-at-a-time is more + # registrations but is the granularity the IB driver expects. + out: List[Tuple[int, int]] = [] + for seg in snapshot: + for block in seg.get("blocks", []) or []: + addr = int(block.get("address", -1)) + size = int(block.get("size", -1)) + state = block.get("state", "") + if addr < 0 or size <= 0 or state != "active_allocated": + continue + if not _overlaps_any(addr, addr + size): + continue + out.append((addr, addr + size)) + return out + + def _merge_against_registered( + self, + candidates: List[Tuple[int, int]], + ) -> List[Tuple[int, int]]: + """Return the (merged) intervals from `candidates` that are not + already covered by any previously-registered range. + + Both sets are merged into one list of disjoint ranges; the diff + against `self._registered_intervals` is the new coverage we need + to ask Mooncake to register. + """ + if not candidates: + return [] + + def merge(intervals: List[Tuple[int, int]]) -> List[Tuple[int, int]]: + sorted_iv = sorted(iv for iv in intervals if iv[1] > iv[0]) + merged: List[Tuple[int, int]] = [] + for s, e in sorted_iv: + if merged and s <= merged[-1][1]: + merged[-1] = (merged[-1][0], max(merged[-1][1], e)) + else: + merged.append((s, e)) + return merged + + cand_merged = merge(candidates) + registered = merge(self._registered_intervals) + new: List[Tuple[int, int]] = [] + + for s, e in cand_merged: + cur_s = s + for rs, re in registered: + if re <= cur_s: + continue + if rs >= e: + break + # rs..re overlaps cur_s..e in some way; carve off the part + # before this registered range. + if rs > cur_s: + new.append((cur_s, rs)) + cur_s = max(cur_s, re) + if cur_s >= e: + break + if cur_s < e: + new.append((cur_s, e)) + return merge(new) + + def _register_persistent_source_intervals( + self, + intervals: List[Tuple[int, int]], + *, + bucket_idx: int, + ) -> None: + """Register stable CUDA source regions once per backend lifetime.""" + if not intervals: + return + if self._engine is None: + raise RuntimeError("[P2P] persistent small registration requires an initialized Mooncake engine") + + # Fast path: if the raw tensor storage interval is already covered by + # a registered active block, skip the expensive CUDA memory snapshot. + raw_new = self._merge_against_registered(intervals) + if not raw_new: + return + + segments = self._intervals_per_cuda_segment(raw_new) + new_segments = self._merge_against_registered(segments) + if not new_segments: + return + + ptrs = [start for start, _ in new_segments] + lengths = [end - start for start, end in new_segments] + ret = self._engine.batch_register(ptrs, lengths) + if ret != 0: + raise RuntimeError( + f"[P2P] persistent small-source batch_register failed: ret={ret} " + f"(bucket {bucket_idx}, regions={len(ptrs)})" + ) + self._registered_source_ptrs.extend(ptrs) + self._registered_intervals.extend(new_segments) + + def _locators_for_source_name(self, name: str) -> Optional[List[Dict[str, Any]]]: + locators = self._tensor_map.get(name) + if locators or name.startswith("language_model."): + return locators + # Kimi-K2.5 SGLang wrappers expose language model tensors under + # language_model.*, while XORL trains the unwrapped text model. + return self._tensor_map.get(f"language_model.{name}") + + def _should_skip_missing_tied_locator(self, name: str) -> bool: + if name != "lm_head.weight": + return False + # Some SGLang models tie lm_head.weight to model.embed_tokens.weight + # and only expose the embedding storage in the P2P tensor map. + # The embedding was already synced from the same root FSDP bucket, so + # failing the whole P2P sync on the missing tied alias is unnecessary. + return bool( + self._tensor_map.get("model.embed_tokens.weight") + or self._tensor_map.get("language_model.model.embed_tokens.weight") + ) + + @staticmethod + def _slice_source_for_locator( + name: str, + full_tensor: torch.Tensor, + loc: Dict[str, Any], + ) -> Optional[torch.Tensor]: + """Extract the sub-region of the trainer's full HF tensor that + corresponds to a single receiver locator. + + The trainer holds the *full* HF tensor (training TP=1 means FSDP + unshard gives the full param). The receiver locator's ``slice`` + field is in full HF coordinates and tells us which rectangle + belongs at this peer's address. + """ + full_tensor = P2PTransportBackend._normalize_source_for_locator(name, full_tensor, loc) + slc = loc.get("slice") + if slc is None: + # Replicated / no sharding — use the whole tensor. + return P2PTransportBackend._normalize_sliced_source_for_locator(name, full_tensor, loc) + + full_shape = loc.get("full_shape") + if full_shape is not None and list(full_tensor.shape) != list(full_shape): + local_view = P2PTransportBackend._slice_qwen35_linear_attention_local_param( + name, + full_tensor, + loc, + full_shape, + slc, + ) + if local_view is not None: + return P2PTransportBackend._normalize_sliced_source_for_locator(name, local_view, loc) + logger.warning( + f"[P2P] full_shape mismatch for {name!r}: " + f"trainer={list(full_tensor.shape)} vs receiver={full_shape}. " + "Check unfuse / quantization in the handler." + ) + return None + + index: Tuple[slice, ...] = tuple(slice(int(s[0]), int(s[1])) for s in slc) + return P2PTransportBackend._normalize_sliced_source_for_locator(name, full_tensor[index], loc) + + @staticmethod + def _normalize_source_for_locator( + name: str, + full_tensor: torch.Tensor, + loc: Dict[str, Any], + ) -> torch.Tensor: + """Normalize trainer-side tensors for receiver-specific HF layouts.""" + full_shape = loc.get("full_shape") + if ( + ".linear_attn." in name + and name.endswith(".conv1d.weight") + and full_shape is not None + and full_tensor.ndim == len(full_shape) + 1 + and full_tensor.shape[1] == 1 + ): + squeezed = full_tensor.squeeze(1) + if list(squeezed.shape) == list(full_shape): + return squeezed + return full_tensor + + @staticmethod + def _normalize_sliced_source_for_locator( + name: str, + local_view: torch.Tensor, + loc: Dict[str, Any], + ) -> torch.Tensor: + """Normalize trainer-side local slices for receiver-specific dtypes.""" + if not ( + ".linear_attn." in name + and (name.endswith(".A_log") or name.endswith(".dt_bias")) + and torch.is_floating_point(local_view) + ): + return local_view + + dtype_name = str(loc.get("dtype") or "") + target_dtype = { + "float32": torch.float32, + "float": torch.float32, + "bfloat16": torch.bfloat16, + "float16": torch.float16, + "half": torch.float16, + }.get(dtype_name) + if target_dtype is None or local_view.dtype == target_dtype: + return local_view + return local_view.to(dtype=target_dtype) + + @staticmethod + def _slice_qwen35_linear_attention_local_param( + name: str, + full_tensor: torch.Tensor, + loc: Dict[str, Any], + full_shape: Any, + slc: Any, + ) -> Optional[torch.Tensor]: + """Handle Qwen3.5 linear-attention locators that expose TP-local vectors. + + Some receiver builds expose ``A_log``/``dt_bias`` locators with the + receiver-local shape (for example ``[8]`` on TP=4) instead of the full + HF shape (``[32]``). The locator still carries ``tp_rank``, so the + sender can recover the intended full-tensor slice. + """ + if not ( + ".linear_attn." in name + and (name.endswith(".A_log") or name.endswith(".dt_bias")) + and full_tensor.ndim == 1 + and isinstance(full_shape, list) + and len(full_shape) == 1 + ): + return None + + local_len = int(full_shape[0]) + if local_len <= 0 or full_tensor.shape[0] <= local_len or full_tensor.shape[0] % local_len != 0: + return None + + if slc is not None: + if len(slc) != 1: + return None + start, stop = int(slc[0][0]), int(slc[0][1]) + if start != 0 or stop != local_len: + return None + + try: + tp_rank = int(loc.get("tp_rank")) + except (TypeError, ValueError): + return None + + tp_size = full_tensor.shape[0] // local_len + if tp_rank < 0 or tp_rank >= tp_size: + return None + + return full_tensor.narrow(0, tp_rank * local_len, local_len) + + def _make_local_engine(self): + """Construct the local Mooncake TransferEngine.""" + try: + # Lazy-import: Mooncake is an optional dep; only pulled when the + # P2P backend is actually selected. + from mooncake.engine import TransferEngine # noqa: PLC0415 + except ImportError as e: + logger.error( + "[P2P] mooncake-transfer-engine is not installed. " + "Install it (see https://kvcache-ai.github.io/Mooncake/getting_started/build.html) " + "or fall back to sync_inference_method='nccl_broadcast'." + ) + logger.error(f"[P2P] underlying ImportError: {e}") + return None + + # Reuse SGLang's Python wrapper so the engine init/handshake is + # identical to the receiver side. We don't depend on SGLang at + # runtime in xorl, but if the package is available locally we use + # it. Otherwise fall back to constructing TransferEngine directly. + hostname = self._hostname or _resolve_local_hostname() + self._resolved_hostname = hostname + logger.info( + "[P2P] local Mooncake endpoint hostname=%s gpu_id=%s ib_device=%s", + hostname, + self._gpu_id, + self._ib_device or "", + ) + try: + from sglang.srt.distributed.device_communicators.mooncake_transfer_engine import ( # noqa: PLC0415 + MooncakeTransferEngine, + ) + + engine = MooncakeTransferEngine( + hostname=hostname, + gpu_id=self._gpu_id, + ib_device=self._ib_device, + ) + return engine + except ImportError: + logger.info( + "[P2P] sglang.srt.distributed.device_communicators.mooncake_transfer_engine " + "is not importable; using mooncake.engine.TransferEngine directly." + ) + try: + return _DirectMooncakeTransferEngine( + TransferEngine, + hostname=hostname, + gpu_id=self._gpu_id, + ib_device=self._ib_device, + ) + except Exception as e: + logger.error(f"[P2P] Failed to initialize direct Mooncake TransferEngine: {e}") + return None + except Exception as e: + logger.error(f"[P2P] Failed to initialize local MooncakeTransferEngine: {e}") + return None + + +def _resolve_local_hostname() -> str: + """Return a routable host:port string for this rank. + + Mooncake's handshake binds on this hostname; it must be reachable from + the SGLang receiver. + """ + # Explicit hostname overrides take precedence (may be FQDN, not just IPv4). + for env_name in ("XORL_P2P_HOSTNAME", "P2P_TRAINER_HOSTNAME", "XORL_WEIGHT_SYNC_MASTER_ADDRESS"): + host = os.environ.get(env_name, "").strip() + if host: + return host + + def _routable_ipv4(value: Optional[str]) -> Optional[str]: + if not value: + return None + try: + addr = ipaddress.ip_address(value.strip()) + except ValueError: + return None + if addr.version != 4 or addr.is_loopback or addr.is_unspecified: + return None + return str(addr) + + for env_name in ("POD_IP", "HOST_IP", "HOSTNAME_IP"): + if ip := _routable_ipv4(os.environ.get(env_name)): + return ip + + # Kubernetes pod hostnames are not necessarily resolvable from peer pods. + # Prefer the local IP address that socket resolves for this pod, matching + # the SGLang receiver's advertised session ids. + try: + if ip := _routable_ipv4(socket.gethostbyname(socket.gethostname())): + return ip + except Exception: + pass + try: + fqdn = socket.getfqdn() + if ip := _routable_ipv4(socket.gethostbyname(fqdn)): + return ip + except Exception: + pass + return socket.gethostname() diff --git a/src/xorl/server/weight_sync/endpoint_manager.py b/src/xorl/server/weight_sync/endpoint_manager.py index 03b31f6b..d98312d0 100644 --- a/src/xorl/server/weight_sync/endpoint_manager.py +++ b/src/xorl/server/weight_sync/endpoint_manager.py @@ -16,6 +16,10 @@ # Reusable session for HTTP connection pooling _http_session: Optional[requests.Session] = None +# SGLang's /health handler can take ~1s even when the server is healthy. +# /model_info is the weight-sync-relevant readiness check and returns quickly; +# keep /health as the last fallback for non-SGLang endpoints. +_HEALTH_PATHS = ("/model_info", "/v1/models", "/health") def _get_http_session() -> requests.Session: @@ -47,13 +51,19 @@ def health_check(self) -> None: """Check all endpoints are healthy. Raises on failure.""" session = _get_http_session() for ep in self.endpoints: - url = f"http://{ep['host']}:{ep['port']}/health" - try: - resp = session.get(url, timeout=10) - 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']} health check failed: {e}") + label = f"{ep['host']}:{ep['port']}" + errors: list[str] = [] + for path in _HEALTH_PATHS: + url = f"http://{label}{path}" + try: + resp = session.get(url, timeout=60) + resp.raise_for_status() + logger.info(f"[EndpointMgr] {label} healthy via {path}") + break + except Exception as e: + errors.append(f"{path}: {e}") + else: + raise RuntimeError(f"Inference endpoint {label} health check failed: {'; '.join(errors)}") def pause( self, diff --git a/src/xorl/server/weight_sync/handler.py b/src/xorl/server/weight_sync/handler.py index b0e1e8c5..8160140b 100644 --- a/src/xorl/server/weight_sync/handler.py +++ b/src/xorl/server/weight_sync/handler.py @@ -26,9 +26,12 @@ This keeps GPU memory to ~1-2 layers (one unsharding + one broadcasting). """ +import atexit import logging +import os import time -from typing import Any, Dict, List, Optional, Tuple +from concurrent.futures import Future, ThreadPoolExecutor +from typing import Any, Callable, Dict, List, Optional, Tuple import torch import torch.distributed as dist @@ -54,8 +57,31 @@ logger = logging.getLogger(__name__) -# Default bucket size for MoE expert broadcasting (256 MB) _DEFAULT_MOE_BUCKET_BYTES = 256 * 1024 * 1024 +_DEFAULT_P2P_MOE_BUCKET_BYTES = 2 * 1024 * 1024 * 1024 + + +def _env_int(name: str, default: int, *, minimum: int = 1) -> int: + raw = os.environ.get(name) + if not raw: + return default + try: + value = int(raw) + except ValueError: + logger.warning("Ignoring invalid %s=%r; using %d", name, raw, default) + return default + if value < minimum: + logger.warning("Ignoring invalid %s=%r; using %d", name, raw, default) + return default + return value + + +def _moe_bucket_size_bytes(sync_method: str) -> int: + """Default MoE bucket sizing is backend-specific; env vars are explicit overrides.""" + default = _DEFAULT_P2P_MOE_BUCKET_BYTES if sync_method == "p2p" else _DEFAULT_MOE_BUCKET_BYTES + if "XORL_WEIGHT_SYNC_MOE_BUCKET_BYTES" in os.environ: + return _env_int("XORL_WEIGHT_SYNC_MOE_BUCKET_BYTES", default) + return _env_int("XORL_WEIGHT_SYNC_BUCKET_BYTES", default) def _prod(shape) -> int: @@ -66,6 +92,262 @@ def _prod(shape) -> int: return r +def _p2p_local_rank(rank: int) -> int: + try: + return int(os.environ.get("LOCAL_RANK", "")) + except ValueError: + if torch.cuda.is_available(): + return rank % max(torch.cuda.device_count(), 1) + return rank + + +def _parse_p2p_gpu_to_ib_map(raw: str) -> Dict[str, str]: + """Parse ``0=mlx5_2,1=mlx5_3`` or ``0:mlx5_2;1:mlx5_3``.""" + mapping: Dict[str, str] = {} + for item in raw.replace(";", ",").split(","): + item = item.strip() + if not item: + continue + sep = "=" if "=" in item else ":" + if sep not in item: + logger.warning("Ignoring malformed P2P_TRAINER_GPU_TO_IB_DEVICE_MAP entry %r", item) + continue + gpu_idx, ib_device = (part.strip() for part in item.split(sep, 1)) + if gpu_idx and ib_device: + mapping[gpu_idx] = ib_device + else: + logger.warning("Ignoring malformed P2P_TRAINER_GPU_TO_IB_DEVICE_MAP entry %r", item) + return mapping + + +def _visible_physical_gpu_indices() -> List[str]: + """Return physical GPU indices in local-rank order when the launcher provides them.""" + for env_name in ("P2P_TRAINER_VISIBLE_GPU_INDICES", "SELECTED_GPU_INDICES"): + raw = os.environ.get(env_name, "").strip() + if raw: + return [idx.strip() for idx in raw.split(",") if idx.strip()] + + # CUDA_VISIBLE_DEVICES is only useful here when it is index-based. In + # Kubernetes it is commonly UUID-based, so callers should provide + # P2P_TRAINER_VISIBLE_GPU_INDICES after their dynamic GPU selection. + raw = os.environ.get("CUDA_VISIBLE_DEVICES", "").strip() + if raw: + indices = [idx.strip() for idx in raw.split(",") if idx.strip()] + if indices and all(idx.isdigit() for idx in indices): + return indices + + return [] + + +def _select_p2p_ib_device(rank: int, world_size: int) -> Optional[str]: + """Return the Mooncake HCA hint for this trainer rank, if configured.""" + per_rank = os.environ.get("P2P_TRAINER_IB_DEVICES_PER_RANK", "").strip() + if per_rank: + devices = [d.strip() for d in per_rank.split(";")] + local_rank = _p2p_local_rank(rank) + if len(devices) >= world_size and 0 <= rank < len(devices): + return devices[rank] or None + if 0 <= local_rank < len(devices): + return devices[local_rank] or None + if 0 <= rank < len(devices): + return devices[rank] or None + logger.warning( + "P2P_TRAINER_IB_DEVICES_PER_RANK has %d entries but no entry for " + "rank=%d local_rank=%d; falling back to P2P_TRAINER_IB_DEVICE/auto-discovery", + len(devices), + rank, + local_rank, + ) + + gpu_to_ib = _parse_p2p_gpu_to_ib_map(os.environ.get("P2P_TRAINER_GPU_TO_IB_DEVICE_MAP", "").strip()) + if gpu_to_ib: + local_rank = _p2p_local_rank(rank) + physical_gpu_indices = _visible_physical_gpu_indices() + physical_gpu_idx = None + if 0 <= local_rank < len(physical_gpu_indices): + physical_gpu_idx = physical_gpu_indices[local_rank] + elif str(local_rank) in gpu_to_ib: + physical_gpu_idx = str(local_rank) + + if physical_gpu_idx is not None: + ib_device = gpu_to_ib.get(physical_gpu_idx) + if ib_device: + return ib_device + logger.warning( + "P2P_TRAINER_GPU_TO_IB_DEVICE_MAP has no entry for physical GPU %s " + "(rank=%d local_rank=%d); falling back to P2P_TRAINER_IB_DEVICE/auto-discovery", + physical_gpu_idx, + rank, + local_rank, + ) + else: + logger.warning( + "P2P_TRAINER_GPU_TO_IB_DEVICE_MAP is set, but local_rank=%d cannot be mapped " + "to a physical GPU index; set P2P_TRAINER_VISIBLE_GPU_INDICES when " + "CUDA_VISIBLE_DEVICES contains GPU UUIDs", + local_rank, + ) + + device = os.environ.get("P2P_TRAINER_IB_DEVICE", "").strip() + return device or None + + +def _safe_abort_token(value: Optional[str]) -> str: + raw = str(value) if value else "none" + safe = "".join(ch if ch.isalnum() or ch in "._-" else "_" for ch in raw) + return safe[:160] or "none" + + +# --------------------------------------------------------------------------- +# Trainer-side P2P backend cache +# +# Constructing the Mooncake TransferEngine + allocating + Mooncake-registering +# the ~8.6 GB of CPU pinned scratch pools costs ~2-3 s on iter 1 (cold) of +# every sync session. Caching the backend across syncs in the same process +# amortizes that cost — second and subsequent syncs reuse the same engine +# and pools and only re-run the per-sync prepare RPC. +# +# Set ``XORL_P2P_BACKEND_CACHE=0`` to disable. +# --------------------------------------------------------------------------- +_cached_p2p_backend: Optional[Any] = None +_cached_backend_key: Optional[Tuple[Any, ...]] = None +_cached_p2p_sender_group: Optional[Any] = None +_cached_p2p_sender_group_ranks: Optional[Tuple[int, ...]] = None + +# nccl_broadcast backend cache — same motivation as P2P cache: the NCCL +# communicator init (TCPStore rendezvous + ncclCommInitRankConfig) is +# expensive and must happen exactly once. Without caching, each call to +# sync_inference_weights creates a new backend and re-sends +# /init_weights_update_group to sglang, which causes sglang to attempt a +# second ncclCommInitRankConfig on the same group name — deadlocking the +# first communicator. Caching ensures initialize() is called only once per +# (endpoint, group_name, world_size) tuple. +_cached_nccl_backend: Optional[Any] = None +_cached_nccl_backend_key: Optional[Tuple[Any, ...]] = None + + +def _atexit_destroy_cached_backend() -> None: + global _cached_p2p_backend, _cached_backend_key, _cached_nccl_backend, _cached_nccl_backend_key + if _cached_p2p_backend is not None: + try: + _cached_p2p_backend.destroy(complete_receiver=False) + except Exception: + pass + _cached_p2p_backend = None + _cached_backend_key = None + if _cached_nccl_backend is not None: + try: + _cached_nccl_backend.destroy(complete_receiver=False) + except Exception: + pass + _cached_nccl_backend = None + _cached_nccl_backend_key = None + + +atexit.register(_atexit_destroy_cached_backend) + + +def _p2p_direct_ep_sender_ranks(ps: Any, world_size: int) -> Tuple[int, ...]: + """Choose one expert sender replica per EP group plus rank 0 for dense weights.""" + if world_size <= 1 or not getattr(ps, "ep_enabled", False): + return tuple(range(world_size)) + ep_mesh = getattr(getattr(ps, "ep_fsdp_device_mesh", None), "mesh", None) + if ep_mesh is None: + return tuple(range(world_size)) + + strategy = os.environ.get("XORL_P2P_DIRECT_EP_REPLICA_STRATEGY", "zero").strip().lower() + if strategy not in {"zero", "round_robin"}: + logger.warning( + "Ignoring invalid XORL_P2P_DIRECT_EP_REPLICA_STRATEGY=%r; using zero", + strategy, + ) + strategy = "zero" + + mesh = ep_mesh.detach().cpu().contiguous() + if mesh.ndim < 2 or mesh.shape[-1] <= 0: + return tuple(range(world_size)) + + ep_size = int(mesh.shape[-2]) + ep_fsdp_size = int(mesh.shape[-1]) + stage_meshes = mesh.reshape(-1, ep_size, ep_fsdp_size) + ranks = {0} + for stage_mesh in stage_meshes: + for ep_rank in range(ep_size): + ep_fsdp_rank = 0 if strategy == "zero" else ep_rank % ep_fsdp_size + rank = int(stage_mesh[ep_rank, ep_fsdp_rank]) + if 0 <= rank < world_size: + ranks.add(rank) + + return tuple(sorted(ranks)) + + +def _p2p_direct_ep_sender_ep_ranks( + ps: Any, + sender_ranks: Tuple[int, ...], + world_size: int, +) -> Tuple[Tuple[int, int], ...]: + """Map each selected sender rank to the EP rank whose experts it owns.""" + if world_size <= 1 or not sender_ranks or not getattr(ps, "ep_enabled", False): + return () + ep_mesh = getattr(getattr(ps, "ep_fsdp_device_mesh", None), "mesh", None) + if ep_mesh is None: + return () + + mesh = ep_mesh.detach().cpu().contiguous() + if mesh.ndim < 2 or mesh.shape[-1] <= 0: + return () + + ep_size = int(mesh.shape[-2]) + ep_fsdp_size = int(mesh.shape[-1]) + stage_meshes = mesh.reshape(-1, ep_size, ep_fsdp_size) + sender_set = set(sender_ranks) + rank_to_ep: Dict[int, int] = {} + for stage_mesh in stage_meshes: + for ep_rank in range(ep_size): + for ep_fsdp_rank in range(ep_fsdp_size): + rank = int(stage_mesh[ep_rank, ep_fsdp_rank]) + if rank in sender_set and rank not in rank_to_ep: + rank_to_ep[rank] = ep_rank + + return tuple((rank, rank_to_ep[rank]) for rank in sender_ranks if rank in rank_to_ep) + + +def _get_p2p_sender_process_group(sender_ranks: Tuple[int, ...], world_size: int) -> Optional[Any]: + """Create/cache a process group for the direct-EP ranks that actually send.""" + if not sender_ranks or sender_ranks == tuple(range(world_size)): + return None + if not dist.is_available() or not dist.is_initialized(): + return None + + global _cached_p2p_sender_group, _cached_p2p_sender_group_ranks + if _cached_p2p_sender_group is not None and _cached_p2p_sender_group_ranks == sender_ranks: + return _cached_p2p_sender_group + + _cached_p2p_sender_group = dist.new_group(ranks=list(sender_ranks)) + _cached_p2p_sender_group_ranks = sender_ranks + return _cached_p2p_sender_group + + +def _backend_cache_value(value: Any) -> Optional[Any]: + if isinstance(value, (str, int, bool, float)): + return value + if isinstance(value, tuple): + cached_items = tuple(_backend_cache_value(item) for item in value) + if all(item is not None for item in cached_items): + return cached_items + return None + + +def _should_collect_ep_moe_tensors(sync_method: str, backend: Any, *, is_sender: bool) -> bool: + """Return whether this rank needs materialized MoE tensors for EP sync.""" + return not ( + sync_method == "p2p" + and getattr(backend, "supports_direct_ep_transfer", False) + and getattr(backend, "has_explicit_sender_ranks", False) + and not is_sender + ) + + class WeightSyncHandler: """Handles weight synchronization between training and inference endpoints.""" @@ -73,6 +355,239 @@ def __init__(self, rank: int, world_size: int, trainer) -> None: self.rank = rank self.world_size = world_size self.trainer = trainer + # Per-sync MoE bucket accumulator. When + # XORL_WEIGHT_SYNC_BATCH_MOE=1, _direct_ep_transfer_experts + # appends here instead of flushing at end-of-call. Caller flushes + # the leftover via _flush_pending_moe_bucket() after the MoE + # loop completes. + self._pending_moe_bucket: List[Tuple[str, torch.Tensor]] = [] + self._pending_moe_bucket_bytes: int = 0 + self._pending_moe_cpu_workspace_records: List[Tuple[str, Tuple[Any, ...], int]] = [] + self._fp8_cpu_workspaces: Dict[Tuple[Any, ...], Dict[str, Any]] = {} + self._p2p_tied_weight_aliases: Dict[str, str] = {} + + def _sync_abort_path(self, group_name: str, weight_version: Optional[str]) -> str: + abort_dir = os.environ.get("XORL_WEIGHT_SYNC_ABORT_DIR", "").strip() + if not abort_dir: + train_config = getattr(self.trainer, "train_config", {}) or {} + if isinstance(train_config, dict): + abort_dir = str(train_config.get("output_dir") or "") + if not abort_dir: + abort_dir = "/tmp" + return os.path.join( + abort_dir, + f".xorl_weight_sync_abort_{_safe_abort_token(group_name)}_{_safe_abort_token(weight_version)}", + ) + + def _clear_sync_abort(self, abort_path: str) -> None: + try: + os.remove(abort_path) + except FileNotFoundError: + pass + except Exception as e: + logger.debug("Rank %d: [WeightSync] failed to clear abort marker %s: %s", self.rank, abort_path, e) + + def _mark_sync_abort(self, abort_path: str, err: Exception) -> None: + try: + abort_dir = os.path.dirname(abort_path) + if abort_dir: + os.makedirs(abort_dir, exist_ok=True) + with open(abort_path, "w", encoding="utf-8") as f: + f.write(f"rank={self.rank}: {err}\n") + except Exception as marker_err: + logger.warning( + "Rank %d: [WeightSync] failed to write abort marker %s: %s", + self.rank, + abort_path, + marker_err, + ) + + def _raise_if_sync_aborted(self, abort_path: str) -> None: + try: + with open(abort_path, encoding="utf-8") as f: + reason = f.read().strip() + except FileNotFoundError: + return + except Exception as e: + logger.debug("Rank %d: [WeightSync] failed to read abort marker %s: %s", self.rank, abort_path, e) + return + + raise RuntimeError(f"Weight sync aborted by peer rank: {reason or abort_path}") + + def _build_p2p_rank_summary( + self, + backend: Any, + *, + is_sender: bool, + transfer_wall_s: float, + total_bytes: int, + num_parameters: int, + num_buckets: int, + ib_device: Optional[str], + phase_s: Dict[str, float], + ) -> Dict[str, Any]: + summary: Dict[str, Any] = { + "rank": self.rank, + "local_rank": _p2p_local_rank(self.rank), + "is_sender": is_sender, + "has_transfers": False, + "transfer_wall_s": transfer_wall_s, + "total_bytes": int(total_bytes), + "num_parameters": int(num_parameters), + "num_buckets": int(num_buckets), + "ib_device": ib_device, + "phase_s": phase_s, + } + if is_sender and hasattr(backend, "stats_summary"): + try: + backend_summary = backend.stats_summary() + summary["backend"] = backend_summary + summary["has_transfers"] = float(backend_summary.get("total_bytes", 0.0)) > 0.0 + backend_main_thread_s = float(backend_summary.get("main_thread_s", 0.0)) + summary["backend_main_thread_s"] = backend_main_thread_s + summary["trainer_overhead_s"] = max( + 0.0, + transfer_wall_s - backend_main_thread_s, + ) + except Exception as e: + summary["backend_error"] = str(e) + return summary + + def _gather_p2p_rank_summaries(self, local_summary: Dict[str, Any]) -> List[Dict[str, Any]]: + if self.world_size <= 1 or not dist.is_available() or not dist.is_initialized(): + return [local_summary] + gathered: List[Any] = [None for _ in range(self.world_size)] + dist.all_gather_object(gathered, local_summary) + return [item for item in gathered if isinstance(item, dict)] + + def _gather_p2p_transfer_statuses(self, local_error: Optional[Exception]) -> List[Dict[str, Any]]: + local_status: Dict[str, Any] = { + "rank": self.rank, + "ok": local_error is None, + } + if local_error is not None: + local_status["error"] = f"{type(local_error).__name__}: {local_error}" + + if self.world_size <= 1 or not dist.is_available() or not dist.is_initialized(): + return [local_status] + gathered: List[Any] = [None for _ in range(self.world_size)] + dist.all_gather_object(gathered, local_status) + return [item for item in gathered if isinstance(item, dict)] + + @staticmethod + def _summary_counter(value: Any) -> int: + try: + return int(value) + except (TypeError, ValueError, OverflowError): + return 0 + + @classmethod + def _aggregate_p2p_transfer_totals( + cls, + p2p_rank_summaries: List[Dict[str, Any]], + *, + total_bytes: int, + num_parameters: int, + num_buckets: int, + ) -> Tuple[int, int, int]: + saw_rank_counters = False + aggregate_bytes = 0 + aggregate_parameters = 0 + aggregate_buckets = 0 + + for summary in p2p_rank_summaries: + if not isinstance(summary, dict): + continue + if not any(key in summary for key in ("total_bytes", "num_parameters", "num_buckets")): + continue + saw_rank_counters = True + aggregate_bytes += cls._summary_counter(summary.get("total_bytes", 0)) + aggregate_parameters += cls._summary_counter(summary.get("num_parameters", 0)) + aggregate_buckets += cls._summary_counter(summary.get("num_buckets", 0)) + + if not saw_rank_counters: + return total_bytes, num_parameters, num_buckets + return aggregate_bytes, aggregate_parameters, aggregate_buckets + + @staticmethod + def _add_rank_timing_breakdown( + timing_breakdown: Dict[str, float], + p2p_rank_summaries: List[Dict[str, Any]], + ) -> None: + sender_transfer_times = [ + float(summary["transfer_wall_s"]) + for summary in p2p_rank_summaries + if summary.get("has_transfers") and isinstance(summary.get("transfer_wall_s"), int | float) + ] + if not sender_transfer_times: + return + max_transfer_s = max(sender_transfer_times) + min_transfer_s = min(sender_transfer_times) + timing_breakdown["max_rank_transfer_s"] = max_transfer_s + timing_breakdown["min_rank_transfer_s"] = min_transfer_s + timing_breakdown["rank_transfer_spread_s"] = max_transfer_s - min_transfer_s + + backend_second_fields = ( + "prepare_s", + "pool_init_s", + "pool_wait_s", + "stage_s", + "submit_s", + "main_thread_s", + "register_s", + "transfer_s", + "deregister_s", + "total_s", + ) + for field in backend_second_fields: + values = [ + float(summary["backend"][field]) + for summary in p2p_rank_summaries + if summary.get("has_transfers") + and isinstance(summary.get("backend"), dict) + and isinstance(summary["backend"].get(field), int | float) + ] + if values: + timing_breakdown[f"p2p_backend_max_{field}"] = max(values) + + @staticmethod + def _moe_runtime_lora_views( + mod: MoEExpertsLoRA | QLoRAMoeExperts, + lora_A: torch.Tensor, + lora_B: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor, float]: + """Return the active LoRA slices and scaling for runtime-rank MoE modules.""" + active_rank = getattr(mod, "active_r", None) + if active_rank is None: + return lora_A, lora_B, float(mod.scaling) + + active_rank = int(active_rank) + if active_rank <= 0: + raise ValueError(f"Active LoRA rank must be positive, got {active_rank}") + if active_rank > lora_A.shape[-1] or active_rank > lora_B.shape[1]: + raise ValueError( + f"Active LoRA rank {active_rank} exceeds available MoE LoRA slices: " + f"A={tuple(lora_A.shape)}, B={tuple(lora_B.shape)}" + ) + + scaling_fn = getattr(mod, "_active_scaling", None) + scaling = float(scaling_fn()) if callable(scaling_fn) else float(mod.scaling) + return lora_A[..., :active_rank], lora_B[:, :active_rank, ...], scaling + + @classmethod + def _compute_moe_lora_delta( + cls, + mod: MoEExpertsLoRA | QLoRAMoeExperts, + lora_A: torch.Tensor, + lora_B: torch.Tensor, + *, + expert_idx: int, + ) -> torch.Tensor: + """Compute one expert's active LoRA delta in GKN format.""" + active_lora_A, active_lora_B, scaling = cls._moe_runtime_lora_views(mod, lora_A, lora_B) + a_idx = min(expert_idx, active_lora_A.shape[0] - 1) + b_idx = min(expert_idx, active_lora_B.shape[0] - 1) + return (active_lora_A[a_idx] @ active_lora_B[b_idx]) * scaling # ======================================================================== # Main entry point @@ -83,7 +598,7 @@ async def handle_sync_inference_weights(self, command_dict: Dict[str, Any]) -> D Handle sync inference weights request (all ranks participate). The ``sync_method`` field selects the transport backend. Currently - supported: ``"nccl_broadcast"``. New backends (RDMA, multi-rank NCCL, + supported: ``"nccl_broadcast"`` and ``"p2p"``. New backends (RDMA, multi-rank NCCL, etc.) can be added by implementing :class:`WeightTransportBackend` and registering in :func:`backends.create_backend`. """ @@ -156,8 +671,13 @@ def _sync_weights( 5. Pipelines: unshard(N+1) overlaps with transfer(N) """ + sync_start_time = time.perf_counter() + timing_breakdown: Dict[str, float] = {} model = self.trainer.model device = f"cuda:{self.trainer.local_rank}" + abort_path = self._sync_abort_path(group_name, weight_version) if sync_method == "p2p" else "" + if abort_path: + self._clear_sync_abort(abort_path) # ------------------------------------------------------------------ # Step 1: Preparation (all ranks) @@ -188,6 +708,58 @@ def _sync_weights( # Step 3: Create backend + endpoint manager, initialize on sender ranks # ------------------------------------------------------------------ + # When sync_method=="p2p" AND EP is active on the trainer, default to + # the multi-sender direct-EP path: each EP rank ships its own local + # experts directly to the receiver instead of dist.gather'ing through + # rank 0 then broadcasting. The default NCCL broadcast backend + # ignores backend_config and stays single-sender. + _ps_for_cfg = get_parallel_state() + _backend_config: Dict[str, Any] = {} + if sync_method == "p2p": + local_rank = _p2p_local_rank(self.rank) + _backend_config["gpu_id"] = local_rank + _backend_config["flush_cache"] = flush_cache + if self._p2p_requires_post_process_weights(quantization): + _backend_config["run_post_process_weights"] = True + if weight_version is not None: + _backend_config["weight_version"] = weight_version + ib_device = _select_p2p_ib_device(self.rank, self.world_size) + if ib_device: + _backend_config["ib_device"] = ib_device + logger.info( + "Rank %d: [WeightSync] P2P Mooncake trainer binding: gpu_id=%s, ib_device=%s", + self.rank, + local_rank, + ib_device or "", + ) + if sync_method == "p2p" and _ps_for_cfg.ep_enabled and _ps_for_cfg.ep_size > 1: + _backend_config["direct_ep_transfer"] = True + # The P2P backend reads world_size out of backend_config; if + # we don't pass it, it defaults to 1 and sender_ranks + # silently collapses to {0} so non-rank-0 trainers route + # back to the gather-and-broadcast fallback. + _backend_config["world_size"] = self.world_size + _backend_config["rank_index"] = self.rank + direct_ep_sender_ranks = _p2p_direct_ep_sender_ranks(_ps_for_cfg, self.world_size) + _backend_config["sender_ranks"] = direct_ep_sender_ranks + _backend_config["direct_ep_size"] = int(_ps_for_cfg.ep_size) + sender_ep_ranks = _p2p_direct_ep_sender_ep_ranks( + _ps_for_cfg, + direct_ep_sender_ranks, + self.world_size, + ) + if sender_ep_ranks: + _backend_config["sender_ep_ranks"] = sender_ep_ranks + sender_group = _get_p2p_sender_process_group(direct_ep_sender_ranks, self.world_size) + if sender_group is not None: + _backend_config["process_group"] = sender_group + logger.info( + "Rank %d: [WeightSync] P2P direct-EP sender ranks=%s sender_ep_ranks=%s", + self.rank, + direct_ep_sender_ranks, + sender_ep_ranks, + ) + transport_cfg = TransportConfig( endpoints=[ EndpointConfig( @@ -204,9 +776,78 @@ def _sync_weights( device=device, training_world_size=self.world_size, training_rank=self.rank, + backend_config=_backend_config, + ) + + # Trainer-side backend cache. The expensive bits — Mooncake + # TransferEngine handshake + ~8.6 GB CPU pinned pool registration — + # are amortized across syncs when (sync_method, endpoint set, + # group_name, master addr) all match the prior call's. The cache + # is module-level so it survives across handler instances within + # the same process. + global _cached_p2p_backend, _cached_backend_key, _cached_nccl_backend, _cached_nccl_backend_key + cache_enabled_p2p = os.environ.get("XORL_P2P_BACKEND_CACHE", "1") == "1" and sync_method == "p2p" + cache_enabled_nccl = os.environ.get("XORL_NCCL_BACKEND_CACHE", "1") == "1" and sync_method in ( + "nccl_broadcast", + "nccl_simple", ) + cache_enabled = cache_enabled_p2p or cache_enabled_nccl + backend_key: Optional[Tuple[Any, ...]] = None + if cache_enabled: + backend_key = ( + sync_method, + tuple((ep["host"], ep["port"], ep.get("world_size", 1)) for ep in endpoints), + group_name, + master_address, + master_port, + buffer_size_mb, + self.world_size, + self.rank, + tuple( + sorted( + (k, cache_v) + for k, v in (_backend_config or {}).items() + if k not in {"flush_cache", "weight_version", "process_group"} + for cache_v in (_backend_cache_value(v),) + if cache_v is not None + ) + ), + ) - backend = create_backend(sync_method, transport_cfg) + # Check nccl_broadcast cache first + if cache_enabled_nccl and _cached_nccl_backend is not None and _cached_nccl_backend_key == backend_key: + backend = _cached_nccl_backend + backend.config = transport_cfg + logger.info(f"Rank {self.rank}: [WeightSync] Reusing cached NCCL backend (skips NCCL group re-init)") + elif ( + cache_enabled_p2p + and _cached_p2p_backend is not None + and _cached_backend_key == backend_key + and getattr(_cached_p2p_backend, "is_alive", False) + ): + backend = _cached_p2p_backend + # Refresh the config in case per-sync params (flush_cache, + # weight_version) differ from the prior call. The cache-key + # check above guarantees the structural fields (endpoints, + # group_name, world_size) match. + backend.config = transport_cfg + logger.info(f"Rank {self.rank}: [WeightSync] Reusing cached P2P backend (skips engine + scratch-pool init)") + else: + if _cached_p2p_backend is not None: + try: + _cached_p2p_backend.destroy(complete_receiver=False) + except Exception as e: + logger.warning(f"[WeightSync] failed to destroy stale cached backend: {e}") + _cached_p2p_backend = None + _cached_backend_key = None + if _cached_nccl_backend is not None: + try: + _cached_nccl_backend.destroy(complete_receiver=False) + except Exception as e: + logger.warning(f"[WeightSync] failed to destroy stale cached nccl backend: {e}") + _cached_nccl_backend = None + _cached_nccl_backend_key = None + backend = create_backend(sync_method, transport_cfg) _is_sender = self.rank in backend.sender_ranks # Endpoint management lives on rank 0 (coordinator). Future multi-rank @@ -216,26 +857,52 @@ def _sync_weights( if self.rank == 0: if not endpoints: return {"success": False, "message": "No endpoints provided"} + t_health = time.perf_counter() endpoint_mgr.health_check() + timing_breakdown["health_check_s"] = time.perf_counter() - t_health # Backend init: all sender ranks participate (collective for NCCL). - if _is_sender: + # For cached backends (nccl_broadcast), initialize() is skipped (already done). + if _is_sender and not ( + (cache_enabled_nccl and _cached_nccl_backend is not None and _cached_nccl_backend_key == backend_key) + or ( + cache_enabled_p2p + and _cached_p2p_backend is not None + and _cached_backend_key == backend_key + and getattr(_cached_p2p_backend, "is_alive", False) + ) + ): logger.info(f"Rank {self.rank}: [WeightSync] Initializing {sync_method} backend...") + t_init = time.perf_counter() if not backend.initialize(): return { "success": False, "message": f"Failed to initialize {sync_method} backend", } + timing_breakdown["backend_init_s"] = time.perf_counter() - t_init logger.info(f"Rank {self.rank}: [WeightSync] Backend initialized") + # Store in cache after successful init (only on rank 0 to avoid races) + if self.rank == 0: + if cache_enabled_nccl: + _cached_nccl_backend = backend + _cached_nccl_backend_key = backend_key + elif cache_enabled_p2p: + _cached_p2p_backend = backend + _cached_backend_key = backend_key # Pause inference: coordinator only (after backend init). if self.rank == 0: logger.info(f"Rank {self.rank}: [WeightSync] Pausing inference (mode={pause_mode})...") + t_pause = time.perf_counter() pause_results, all_paused = endpoint_mgr.pause(pause_mode) + timing_breakdown["pause_s"] = time.perf_counter() - t_pause if not all_paused: endpoint_mgr.resume() if _is_sender: - backend.destroy() + backend.destroy(complete_receiver=False) + # Pause failure invalidates the cache; next sync starts fresh. + _cached_p2p_backend = None + _cached_backend_key = None return { "success": False, "message": f"Failed to pause inference endpoints: {pause_results}", @@ -248,6 +915,24 @@ def _sync_weights( total_bytes = 0 total_params = 0 num_buckets = 0 + rank_phase_s: Dict[str, float] = {} + + def _add_rank_phase(name: str, start: float) -> None: + rank_phase_s[name] = rank_phase_s.get(name, 0.0) + (time.perf_counter() - start) + + # Cross-layer MoE batching. When on, _direct_ep_transfer_experts + # appends to the handler-level _pending_moe_bucket instead of flushing + # at end-of-call; we ship the leftover once after the module loop. + # For the scaled P2P path, pair this with an explicit + # XORL_WEIGHT_SYNC_MOE_BUCKET_BYTES so MoE batching does not silently + # reuse the dense/root bucket cap. Default off for non-P2P/back-compat. + batch_moe = os.environ.get("XORL_WEIGHT_SYNC_BATCH_MOE", "0") == "1" + moe_bucket_size_bytes = _moe_bucket_size_bytes(sync_method) + # Reset cross-sync state in case a prior sync raised mid-flush. + self._pending_moe_bucket = [] + self._pending_moe_bucket_bytes = 0 + self._p2p_tied_weight_aliases = {} + self._reset_fp8_cpu_workspace_usage() # Build ordered list of FSDP modules to process modules_to_sync: List[Tuple[str, FSDPModule]] = [] @@ -258,6 +943,21 @@ def _sync_weights( # Detect EP mode _ps = get_parallel_state() _ep_enabled = _ps.ep_enabled and _ps.ep_size > 1 + _collect_ep_moe_tensors = _should_collect_ep_moe_tensors( + sync_method, + backend, + is_sender=_is_sender, + ) + _extract_dense_on_sender = bool( + sync_method == "p2p" + and _is_sender + and getattr(backend, "should_extract_dense_params_on_rank", lambda _rank: False)(self.rank) + ) + if _ep_enabled and not _collect_ep_moe_tensors: + logger.info( + "Rank %d: [WeightSync] Skipping EP MoE tensor materialization on non-sender P2P rank", + self.rank, + ) # Detect PP mode (Pipeline Parallelism) # With PP, each stage has an independent FSDP shard group. We process @@ -295,31 +995,84 @@ def _sync_weights( f"({num_stage_modules} modules, remote={_is_remote})" ) + # Optional fast path: unshard ALL FSDP modules up front so + # the per-module loop doesn't pay the FSDP allgather barrier + # latency (~50-100 ms × 50 modules = 2.5-5 s of barrier time + # collapsed to one batched pass). Memory cost: each rank + # holds the full model in addition to the sharded copy + # (~30 GB extra on Qwen3-30B-A3B at FSDP=8). Gate behind + # XORL_WEIGHT_SYNC_PRE_UNSHARD=1; off by default. + _pre_unshard = os.environ.get("XORL_WEIGHT_SYNC_PRE_UNSHARD", "0") == "1" and _is_my_stage + if _pre_unshard: + t_pre = time.perf_counter() + for _, _fsdp_mod in stage_modules: + _fsdp_mod.unshard() + # No torch.cuda.synchronize() — unshards queue on + # the NCCL stream and the first GPU op in + # _extract_params_for_sync will naturally wait via + # stream ordering. Skipping the sync lets the + # streaming loop start ~1-2 s earlier on rank 0 + # (which had ~2s of launch latency relative to + # other ranks in baseline measurements). + logger.info( + f"Rank {self.rank}: [WeightSync] Pre-unshard launch done: " + f"{len(stage_modules)} modules queued in " + f"{(time.perf_counter() - t_pre) * 1000:.1f} ms " + f"(allocated={torch.cuda.memory_allocated() / 1e9:.2f} GB)" + ) + + # XORL_WEIGHT_SYNC_TIMINGS=1 → emit a per-module phase + # breakdown on rank 0 (unshard / qlora / ep_collect / + # extract / unfuse / broadcast / direct_ep). Pinpoints + # which trainer-side phase dominates the streaming wall. + _ws_timings = os.environ.get("XORL_WEIGHT_SYNC_TIMINGS", "0") == "1" + for mod_idx in range(num_stage_modules): + if abort_path: + self._raise_if_sync_aborted(abort_path) 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 moe_contexts = [] ep_moe_contexts = [] + _t0 = time.perf_counter() if _ws_timings else 0.0 + _t_unshard = _t_qlora = _t_ep_collect = _t_extract = _t0 + _t_unfuse = _t_broadcast = _t_direct_ep = _t0 if _is_my_stage: mod_name, fsdp_mod = stage_modules[mod_idx] - fsdp_mod.unshard() + if not _pre_unshard: + t_phase = time.perf_counter() + fsdp_mod.unshard() + _add_rank_phase("unshard_s", t_phase) + if _ws_timings: + _t_unshard = time.perf_counter() + t_phase = time.perf_counter() qlora_linear_buffer, moe_contexts = self._qlora_collective_ops( fsdp_mod, mod_name, collect_results=_stage_leader, ) + _add_rank_phase("qlora_s", t_phase) + if _ws_timings: + _t_qlora = time.perf_counter() if _ep_enabled: + t_phase = time.perf_counter() ep_moe_contexts = self._collect_ep_moe_data( fsdp_mod, mod_name, _ps, + skip_clone=_pre_unshard, + collect_tensors=_collect_ep_moe_tensors, + phase_s=rank_phase_s, ) + _add_rank_phase("ep_collect_s", t_phase) + if _ws_timings: + _t_ep_collect = time.perf_counter() # EP MoE prefixes to skip in extraction ep_moe_prefixes = set() @@ -333,74 +1086,176 @@ def _sync_weights( else: ep_moe_prefixes.add(p) - if _stage_leader: + if _stage_leader or _extract_dense_on_sender: + t_phase = time.perf_counter() if ep_moe_prefixes: logger.info( f"Rank {self.rank}: [WeightSync] ep_moe_prefixes={ep_moe_prefixes} for {mod_name}" ) + include_dense_param: Optional[Callable[[str], bool]] = None + if _extract_dense_on_sender: + should_send_dense_param = getattr( + backend, + "should_send_dense_param", + lambda _name, _rank: True, + ) + + def include_dense_param(name: str, _rank: int = self.rank) -> bool: + return bool(should_send_dense_param(name, _rank)) + current_buffer = self._extract_params_for_sync( fsdp_mod, mod_name, DTensor, skip_moe_prefixes=ep_moe_prefixes, + emit_tied_weight_duplicates=not ( + sync_method == "p2p" + and os.environ.get("XORL_P2P_TIED_WEIGHT_ALIAS_COPY", "1") == "1" + ), + tied_weight_aliases=self._p2p_tied_weight_aliases, + include_param=include_dense_param, ) current_buffer.extend(qlora_linear_buffer) + _add_rank_phase("extract_s", t_phase) del qlora_linear_buffer + if _ws_timings: + _t_extract = time.perf_counter() - fsdp_mod.reshard() + if not _pre_unshard: + t_phase = time.perf_counter() + fsdp_mod.reshard() + _add_rank_phase("reshard_s", t_phase) # ── Transfer / broadcast to SGLang ─────────────────── if not _is_remote: # Stage 0: sender rank(s) broadcast directly to SGLang if _is_sender and current_buffer: - current_buffer = self._unfuse_for_inference( - current_buffer, - model, - ) - if quantization and quantization.get("quant_method") == "fp8": - current_buffer = self._quantize_buffer_for_fp8( + t_phase = time.perf_counter() + if current_buffer: + current_buffer = self._unfuse_for_inference( current_buffer, - quantization_config=quantization, + model, ) - 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), - weight_version=weight_version if is_last_overall and not moe_contexts else None, - ) - total_bytes += b - total_params += p - num_buckets += 1 - del current_buffer - - # Stage 0 MoE handling (unchanged) + current_buffer = getattr( + backend, + "filter_dense_buffer_for_rank", + lambda buf, _rank: buf, + )(current_buffer, self.rank) + if current_buffer: + if quantization and quantization.get("quant_method") == "fp8": + current_buffer = self._quantize_buffer_for_fp8( + current_buffer, + quantization_config=quantization, + target_device=self._fp8_quantization_target_device(backend), + phase_s=rank_phase_s, + phase_prefix="dense_fp8", + ) + else: + logger.debug( + "Rank %d: [WeightSync] Dense shard filter skipped module %s", + self.rank, + mod_name, + ) + _add_rank_phase("unfuse_quantize_s", t_phase) + if _ws_timings: + _t_unfuse = time.perf_counter() + if current_buffer: + logger.info( + f"Rank {self.rank}: [WeightSync] Module {mod_name}: {len(current_buffer)} params" + ) + t_phase = time.perf_counter() + b, p = self._broadcast_buffer( + 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, + ) + _add_rank_phase("broadcast_buffer_s", t_phase) + total_bytes += b + total_params += p + num_buckets += 1 + del current_buffer + if _ws_timings: + _t_broadcast = time.perf_counter() + + # Stage 0 MoE handling. With direct EP/PP transport + # (P2P + direct_ep_transfer=True), each EP rank ships + # its own local experts in parallel and skips the + # rank-0 dist.gather → broadcast funnel. The default + # NCCL path still does gather-and-broadcast. if moe_contexts or ep_moe_contexts: if _ep_enabled: + use_direct_ep = backend.supports_direct_ep_transfer for ctx in moe_contexts + ep_moe_contexts: - b, p, n = self._gather_and_broadcast_ep_moe_experts( - backend, - ctx, - flush_cache=(flush_cache and is_last_overall), - weight_version=weight_version if is_last_overall else None, - quantization=quantization, - ps=_ps, - ) + if use_direct_ep: + # batch_moe defers the per-call + # final flush so multiple layers' + # MoE experts coalesce into one + # large bucket (~2 GB instead of + # ~302 MB). flush_cache and + # weight_version migrate to the + # post-loop _flush_pending_moe_bucket + # call below. + t_phase = time.perf_counter() + b, p, n = self._direct_ep_transfer_experts( + backend, + ctx, + flush_cache=(flush_cache and is_last_overall) and not batch_moe, + weight_version=( + weight_version if is_last_overall and not batch_moe else None + ), + bucket_size_bytes=moe_bucket_size_bytes, + quantization=quantization, + ps=_ps, + defer_final_flush=batch_moe, + phase_s=rank_phase_s, + ) + _add_rank_phase("direct_ep_s", t_phase) + else: + t_phase = time.perf_counter() + b, p, n = self._gather_and_broadcast_ep_moe_experts( + backend, + ctx, + flush_cache=(flush_cache and is_last_overall), + weight_version=weight_version if is_last_overall else None, + bucket_size_bytes=moe_bucket_size_bytes, + quantization=quantization, + ps=_ps, + ) + _add_rank_phase("gather_broadcast_ep_s", t_phase) total_bytes += b total_params += p num_buckets += n elif _is_sender: for ctx in moe_contexts: + t_phase = time.perf_counter() b, p, n = self._broadcast_moe_experts_bucketed( backend, ctx, flush_cache=(flush_cache and is_last_overall), weight_version=weight_version if is_last_overall else None, + bucket_size_bytes=moe_bucket_size_bytes, quantization=quantization, ) + _add_rank_phase("broadcast_moe_s", t_phase) total_bytes += b total_params += p num_buckets += n + + if _ws_timings and _is_my_stage and self.rank == 0: + _t_direct_ep = time.perf_counter() + _mn = stage_modules[mod_idx][0] if mod_idx < len(stage_modules) else "?" + logger.info( + f"Rank 0: [WeightSync timing] {_mn}: " + f"unshard={(_t_unshard - _t0) * 1000:.0f}ms " + f"qlora={(_t_qlora - _t_unshard) * 1000:.0f}ms " + f"ep_collect={(_t_ep_collect - _t_qlora) * 1000:.0f}ms " + f"extract={(_t_extract - _t_ep_collect) * 1000:.0f}ms " + f"unfuse={(_t_unfuse - _t_extract) * 1000:.0f}ms " + f"broadcast={(_t_broadcast - _t_unfuse) * 1000:.0f}ms " + f"direct_ep={(_t_direct_ep - _t_broadcast) * 1000:.0f}ms " + f"total={(_t_direct_ep - _t0) * 1000:.0f}ms" + ) else: # Remote stage: per-module NCCL transfer to rank 0 if _ps.dp_shard_rank == 0: @@ -428,6 +1283,9 @@ def _sync_weights( received = self._quantize_buffer_for_fp8( received, quantization_config=quantization, + target_device=self._fp8_quantization_target_device(backend), + phase_s=rank_phase_s, + phase_prefix="pp_fp8", ) logger.info( f"Rank 0: [WeightSync] PP stage {pp_stage} module " @@ -444,15 +1302,150 @@ def _sync_weights( num_buckets += 1 del received + if abort_path: + self._raise_if_sync_aborted(abort_path) + + # Pre-unshard mode: now that all transfers have been + # submitted to the worker (transfer_bucket returns after + # staging), re-shard the modules to free the ~30 GB of + # extra GPU memory that's been holding the unsharded + # weights. We do this BEFORE flush_pending_transfers so + # the reshard work can happen on the compute stream + # while RDMA reads from the CPU pinned pool on the NIC. + if _pre_unshard: + t_re = time.perf_counter() + for _, _fsdp_mod in stage_modules: + _fsdp_mod.reshard() + logger.info( + f"Rank {self.rank}: [WeightSync] Post-streaming reshard " + f"in {(time.perf_counter() - t_re) * 1000:.1f} ms" + ) + # Barrier between PP stages (all ranks) if _pp_enabled: dist.barrier() + # Cross-layer MoE flush. Ship whatever's left in the + # accumulator once, instead of per-layer. This is the LAST + # transfer of the sync, so it carries flush_cache + + # weight_version (if requested by the caller). + if abort_path: + self._raise_if_sync_aborted(abort_path) + if batch_moe and _is_sender: + t_phase = time.perf_counter() + b, p, n = self._flush_pending_moe_bucket( + backend, + flush_cache=flush_cache, + weight_version=weight_version, + quantization=quantization, + bucket_size_bytes=moe_bucket_size_bytes, + phase_s=rank_phase_s, + ) + _add_rank_phase("moe_final_flush_s", t_phase) + total_bytes += b + total_params += p + num_buckets += n + + # Drain any async transfers (P2P backend submits Mooncake + # work to a worker thread and returns from transfer_bucket + # before bytes land). Must complete before the handler + # resumes inference or the next request can read + # partially-updated weights. + pending_transfer_error: Optional[Exception] = None + if abort_path: + try: + self._raise_if_sync_aborted(abort_path) + except Exception as abort_err: + pending_transfer_error = abort_err + if _is_sender: + t_phase = time.perf_counter() + try: + if pending_transfer_error is None: + backend.flush_pending_transfers() + except Exception as flush_err: + pending_transfer_error = flush_err + if abort_path: + self._mark_sync_abort(abort_path, flush_err) + finally: + _add_rank_phase("flush_pending_s", t_phase) + + if sync_method == "p2p": + transfer_statuses = self._gather_p2p_transfer_statuses(pending_transfer_error) + failed_statuses = [status for status in transfer_statuses if not status.get("ok", False)] + if failed_statuses: + if pending_transfer_error is not None: + raise pending_transfer_error + preview = "; ".join( + f"rank {status.get('rank')}: {status.get('error', 'unknown error')}" + for status in failed_statuses[:4] + ) + if len(failed_statuses) > 4: + preview += f"; ... {len(failed_statuses) - 4} more" + raise RuntimeError(f"P2P transfer failed on peer rank(s): {preview}") + elif pending_transfer_error is not None: + raise pending_transfer_error + transfer_time = time.perf_counter() - start_time + timing_breakdown["transfer_s"] = transfer_time + p2p_rank_summaries: List[Dict[str, Any]] = [] + if sync_method == "p2p": + t_rank_summary = time.perf_counter() + local_summary = self._build_p2p_rank_summary( + backend, + is_sender=_is_sender, + transfer_wall_s=transfer_time, + total_bytes=total_bytes, + num_parameters=total_params, + num_buckets=num_buckets, + ib_device=_backend_config.get("ib_device"), + phase_s=rank_phase_s, + ) + p2p_rank_summaries = self._gather_p2p_rank_summaries(local_summary) + total_bytes, total_params, num_buckets = self._aggregate_p2p_transfer_totals( + p2p_rank_summaries, + total_bytes=total_bytes, + num_parameters=total_params, + num_buckets=num_buckets, + ) + if self.rank == 0: + timing_breakdown["rank_summary_gather_s"] = time.perf_counter() - t_rank_summary + self._add_rank_timing_breakdown(timing_breakdown, p2p_rank_summaries) # ------------------------------------------------------------------ # Step 5: Resume inference, cleanup # ------------------------------------------------------------------ + if _is_sender: + if sync_method == "p2p" and self._p2p_tied_weight_aliases: + backend.config.backend_config["p2p_tied_weight_aliases"] = dict(self._p2p_tied_weight_aliases) + # Finalize receiver-side update before inference resumes. + # For P2P this sends /complete_weights_update, where SGLang + # applies weight_version, flush_cache, and post-processing. + # If completion fails, fail closed and leave inference paused. + if cache_enabled and backend_key is not None and hasattr(backend, "complete_sync"): + t_complete = time.perf_counter() + try: + backend.complete_sync() + _cached_p2p_backend = backend + _cached_backend_key = backend_key + except Exception as complete_err: + logger.warning( + f"Rank {self.rank}: [WeightSync] complete_sync failed; " + f"falling back to full destroy: {complete_err}" + ) + try: + backend.destroy(complete_receiver=False) + except Exception: + pass + _cached_p2p_backend = None + _cached_backend_key = None + raise + finally: + timing_breakdown["complete_s"] = time.perf_counter() - t_complete + else: + t_destroy = time.perf_counter() + backend.destroy() + timing_breakdown["backend_destroy_s"] = time.perf_counter() - t_destroy + if self.rank == 0: throughput = (total_bytes / transfer_time / (1024**3)) if transfer_time > 0 else 0 logger.info( @@ -461,33 +1454,52 @@ def _sync_weights( f"{total_bytes / 1e9:.2f} GB, {total_params} params, " f"{num_buckets} buckets" ) + t_resume = time.perf_counter() endpoint_mgr.resume() - if _is_sender: - backend.destroy() + timing_breakdown["resume_s"] = time.perf_counter() - t_resume + + timing_breakdown["total_handler_s"] = time.perf_counter() - sync_start_time + if self.rank == 0: + ordered = ", ".join(f"{k}={v:.3f}s" for k, v in timing_breakdown.items()) + logger.info(f"Rank {self.rank}: [WeightSync] Timing breakdown: {ordered}") return { "success": True, "message": f"Synced {total_params} params to {len(endpoints)} endpoint(s)", - "transfer_time": time.perf_counter() - start_time, + "transfer_time": transfer_time, "total_bytes": total_bytes, "num_parameters": total_params, "num_buckets": num_buckets, + "timing_breakdown": timing_breakdown, + "p2p_rank_summaries": p2p_rank_summaries, "endpoint_results": [{"host": ep["host"], "port": ep["port"], "success": True} for ep in endpoints], } - except Exception: + except Exception as sync_err: + if abort_path: + self._mark_sync_abort(abort_path, sync_err) if endpoint_mgr is not None: - try: - endpoint_mgr.resume() - except Exception as resume_err: - logger.warning(f"Rank 0: [WeightSync] Failed to resume inference during cleanup: {resume_err}") + if sync_method == "p2p": + logger.warning( + "Rank 0: [WeightSync] P2P sync failed after streaming began; " + "not resuming inference because RDMA may have partially updated receiver weights" + ) + else: + try: + endpoint_mgr.resume() + except Exception as resume_err: + logger.warning(f"Rank 0: [WeightSync] Failed to resume inference during cleanup: {resume_err}") if _is_sender: try: - backend.destroy() + backend.destroy(complete_receiver=False) except Exception as destroy_err: logger.warning( f"Rank {self.rank}: [WeightSync] Failed to destroy backend during cleanup: {destroy_err}" ) + # Failure path always invalidates the cache so a fresh + # backend is created next sync. + _cached_p2p_backend = None + _cached_backend_key = None raise # ======================================================================== @@ -620,12 +1632,17 @@ def _collect_ep_moe_data( fsdp_mod, mod_name: str, ps, + skip_clone: bool = False, + collect_tensors: bool = True, + phase_s: Optional[Dict[str, float]] = None, ) -> List[Dict[str, Any]]: """Collect local EP-sharded MoE expert data during unshard phase. Identifies full-weight MoE modules (MoEExperts, MoEExpertsLoRA) whose expert params are EP-sharded DTensors. Clones local expert data for - later EP gathering after reshard. + later EP gathering after reshard when ``collect_tensors`` is true. + Non-sender direct-P2P ranks can set ``collect_tensors=False`` to keep + MoE prefix metadata without cloning tensors they will not transfer. QLoRAMoeExperts are handled separately by _qlora_collective_ops. @@ -665,6 +1682,17 @@ def _collect_ep_moe_data( else: full_prefix = mod_name + if not collect_tensors: + contexts.append( + { + "type": "full_weight", + "prefix": full_prefix, + "local_experts": None, + "num_local_experts": E_local, + } + ) + continue + # Clone local expert data for each projection. # With EP, each rank's module already holds only local experts [E_local, K, N]. local_experts = {} @@ -678,16 +1706,30 @@ def _collect_ep_moe_data( # Merge LoRA if applicable if isinstance(mod, MoEExpertsLoRA): if proj_name in mod.lora_config.target_modules: + t_convert = time.perf_counter() delta = mod._compute_proj_delta(proj_name) if isinstance(delta, DTensor): delta = delta.to_local() local = local.to(torch.bfloat16) + delta.to(torch.bfloat16) + self._add_phase_time(phase_s, "ep_collect_convert_s", time.perf_counter() - t_convert) else: + t_convert = time.perf_counter() local = local.to(torch.bfloat16) + self._add_phase_time(phase_s, "ep_collect_convert_s", time.perf_counter() - t_convert) else: + t_convert = time.perf_counter() local = local.to(torch.bfloat16) - - local_experts[proj_name] = local.clone() + self._add_phase_time(phase_s, "ep_collect_convert_s", time.perf_counter() - t_convert) + + # Pre-unshard mode: the unsharded module storage stays + # alive across the whole streaming loop (we reshard + # everything at the end), so we can hand out a view + # instead of cloning. With pre-unshard off, .clone() is + # required because the per-iteration reshard will free + # the source memory before transfer reads it. + t_clone = time.perf_counter() + local_experts[proj_name] = local if skip_clone else local.clone() + self._add_phase_time(phase_s, "ep_collect_clone_s", time.perf_counter() - t_clone) contexts.append( { @@ -704,6 +1746,512 @@ def _collect_ep_moe_data( # EP-aware MoE expert gathering and broadcasting (all ranks) # ======================================================================== + def _direct_ep_transfer_experts( + self, + 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, + defer_final_flush: bool = False, + phase_s: Optional[Dict[str, float]] = None, + ) -> Tuple[int, int, int]: + """Multi-sender EP path: each rank ships its own local experts. + + Compared to :meth:`_gather_and_broadcast_ep_moe_experts`, this + skips the per-projection ``dist.gather → rank 0 → broadcast`` + funnel. Each EP rank formats its own ``ctx["local_experts"]`` + as HF-named per-expert tensors and calls + ``backend.transfer_bucket(..., src_rank=self.rank)``. With N EP + ranks, aggregate trainer→inference bandwidth scales N×. + + Falls back to the gather path for QLoRA contexts — the per-rank + lora-merge path is similar in shape but model-specific and + tracked as a follow-up. + + Like the gather path, only the EP-FSDP-rank-0 replica column + sends; other replicas have identical local shards and would + duplicate data on the wire. + """ + # Backend must declare direct-EP support; fall back if not. + if not backend.supports_direct_ep_transfer: + return self._gather_and_broadcast_ep_moe_experts( + backend, + ctx, + flush_cache=flush_cache, + weight_version=weight_version, + bucket_size_bytes=bucket_size_bytes, + quantization=quantization, + ps=ps, + ) + + is_qlora = ctx.get("type") != "full_weight" + if is_qlora: + # QLoRA direct-EP needs per-rank dequantize + lora merge into + # an HF-shaped buffer, mirroring the gather path's lora math + # but without the gather. Tracked as follow-up; defer to the + # gather implementation today so QLoRA users still ship. + return self._gather_and_broadcast_ep_moe_experts( + backend, + ctx, + flush_cache=flush_cache, + weight_version=weight_version, + bucket_size_bytes=bucket_size_bytes, + quantization=quantization, + ps=ps, + ) + + if getattr(backend, "has_explicit_sender_ranks", False): + if self.rank not in backend.sender_ranks: + ctx["local_experts"] = None + return 0, 0, 0 + else: + 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") + if ep_fsdp_rank != 0: + ctx["local_experts"] = None + return 0, 0, 0 + + full_prefix = ctx["prefix"] + ep_size = ps.ep_size + ep_rank = ps.ep_rank + local_experts = ctx["local_experts"] + E_local = ctx["num_local_experts"] + + logger.info( + f"Rank {self.rank}: [Direct-EP] prefix={full_prefix}, E_local={E_local}, E_total={E_local * ep_size}" + ) + + total_bytes = 0 + total_params = 0 + num_buckets = 0 + fp8_cpu_workspace_pending_source_limit = self._fp8_cpu_workspace_pending_source_bytes(bucket_size_bytes) + + def flush_bucket_before_append(next_entry_bytes: int) -> None: + nonlocal bucket, bucket_bytes, num_buckets + if not self._would_exceed_bucket_cap(bucket_bytes, next_entry_bytes, bucket_size_bytes): + return + t_backend = time.perf_counter() + backend.transfer_bucket( + bucket, + src_rank=self.rank, + flush_cache=False, + ) + self._add_phase_time(phase_s, "direct_ep_backend_s", time.perf_counter() - t_backend) + bucket = [] + bucket_bytes = 0 + num_buckets += 1 + + # When batch mode defers the final flush, append to the handler-level + # bucket so later MoE calls can coalesce into the same transfer. + if defer_final_flush: + bucket = self._pending_moe_bucket + bucket_bytes = self._pending_moe_bucket_bytes + else: + bucket = [] + bucket_bytes = 0 + device = f"cuda:{self.rank % torch.cuda.device_count()}" + fp8_cpu_quantization = ( + quantization is not None + and quantization.get("quant_method") == "fp8" + and self._fp8_quantization_target_device(backend) == "cpu" + ) + fp8_gpu_quantization = ( + fp8_cpu_quantization + and self._fp8_quantization_execution_device() in {"gpu", "cuda"} + and quantization.get("fmt", "e4m3") == "e4m3" + ) + fp8_cpu_workspace = ( + fp8_cpu_quantization + and not fp8_gpu_quantization + and defer_final_flush + and self._fp8_cpu_workspace_enabled() + and not quantization.get("modules_to_not_convert") + ) + + # local_experts[proj] is [E_local, K, N] (input-major). HF + # convention is [N, K] per-expert (output-major) — same permute + # the gather path does before broadcast. + for proj_name in ("gate_proj", "up_proj", "down_proj"): + logger.debug(f"Rank {self.rank}: [Direct-EP] {full_prefix}.{proj_name} stage=before_permute") + local_data = local_experts[proj_name] # [E_local, K, N] + if fp8_gpu_quantization and local_data.device.type == "cuda": + entries, original_bytes = self._quantize_ep_expert_projection_for_fp8_gpu_to_cpu( + local_data, + full_prefix=full_prefix, + proj_name=proj_name, + ep_rank=ep_rank, + quantization_config=quantization, + phase_s=phase_s, + ) + total_bytes += original_bytes + total_params += E_local + for entry_name, entry_tensor in entries: + entry_bytes = entry_tensor.numel() * entry_tensor.element_size() + flush_bucket_before_append(entry_bytes) + bucket.append((entry_name, entry_tensor)) + bucket_bytes += entry_bytes + + if bucket_bytes >= bucket_size_bytes: + t_backend = time.perf_counter() + backend.transfer_bucket( + bucket, + src_rank=self.rank, + flush_cache=False, + ) + self._add_phase_time(phase_s, "direct_ep_backend_s", time.perf_counter() - t_backend) + bucket = [] + bucket_bytes = 0 + num_buckets += 1 + continue + + if fp8_cpu_workspace: + records, original_bytes = self._stage_ep_expert_projection_for_fp8_cpu_workspace( + local_data, + full_prefix=full_prefix, + proj_name=proj_name, + ep_rank=ep_rank, + quantization_config=quantization, + phase_s=phase_s, + ) + total_bytes += original_bytes + total_params += E_local + self._pending_moe_cpu_workspace_records.extend(records) + bucket_bytes += original_bytes + self._pending_moe_bucket_bytes = bucket_bytes + if bucket_bytes >= fp8_cpu_workspace_pending_source_limit: + _, _, flushed_buckets = self._flush_pending_moe_bucket( + backend, + flush_cache=False, + weight_version=None, + quantization=quantization, + bucket_size_bytes=bucket_size_bytes, + phase_s=phase_s, + ) + num_buckets += flushed_buckets + bucket = self._pending_moe_bucket + bucket_bytes = self._pending_moe_bucket_bytes + continue + + if fp8_cpu_quantization: + entries, original_bytes = self._quantize_ep_expert_projection_for_fp8_cpu( + local_data, + full_prefix=full_prefix, + proj_name=proj_name, + ep_rank=ep_rank, + quantization_config=quantization, + phase_s=phase_s, + ) + total_bytes += original_bytes + total_params += E_local + for entry_name, entry_tensor in entries: + entry_bytes = entry_tensor.numel() * entry_tensor.element_size() + flush_bucket_before_append(entry_bytes) + bucket.append((entry_name, entry_tensor)) + bucket_bytes += entry_bytes + + if bucket_bytes >= bucket_size_bytes: + t_backend = time.perf_counter() + backend.transfer_bucket( + bucket, + src_rank=self.rank, + flush_cache=False, + ) + self._add_phase_time(phase_s, "direct_ep_backend_s", time.perf_counter() - t_backend) + bucket = [] + bucket_bytes = 0 + num_buckets += 1 + continue + + t_permute = time.perf_counter() + local_stack = local_data.permute(0, 2, 1).contiguous().to(device) + self._add_phase_time(phase_s, "direct_ep_permute_s", time.perf_counter() - t_permute) + logger.debug( + f"Rank {self.rank}: [Direct-EP] {full_prefix}.{proj_name} " + f"stage=after_permute shape={tuple(local_stack.shape)} dtype={local_stack.dtype}" + ) + for i in range(E_local): + global_idx = ep_rank * E_local + i + hf_name = f"{full_prefix}.{global_idx}.{proj_name}.weight" + tensor = local_stack[i] + tensor_bytes = tensor.numel() * tensor.element_size() + flush_bucket_before_append(tensor_bytes) + bucket.append((hf_name, tensor)) + bucket_bytes += tensor_bytes + total_bytes += tensor_bytes + total_params += 1 + + if bucket_bytes >= bucket_size_bytes: + if quantization and quantization.get("quant_method") == "fp8": + bucket = self._quantize_buffer_for_fp8( + bucket, + quantization_config=quantization, + target_device=self._fp8_quantization_target_device(backend), + phase_s=phase_s, + phase_prefix="direct_ep_fp8", + ) + t_backend = time.perf_counter() + backend.transfer_bucket( + bucket, + src_rank=self.rank, + flush_cache=False, + ) + self._add_phase_time(phase_s, "direct_ep_backend_s", time.perf_counter() - t_backend) + bucket = [] + bucket_bytes = 0 + num_buckets += 1 + del local_stack + + if defer_final_flush: + # Hand the partial bucket back to the handler-level state + # so the next layer's MoE call (or the final flush) picks + # it up. Skip the per-call final flush entirely. + self._pending_moe_bucket = bucket + self._pending_moe_bucket_bytes = bucket_bytes + logger.debug( + f"Rank {self.rank}: [Direct-EP] {full_prefix} " + f"stage=defer_final_flush bucket_bytes={bucket_bytes} bucket_len={len(bucket)}" + ) + elif bucket: + logger.debug( + f"Rank {self.rank}: [Direct-EP] {full_prefix} " + f"stage=before_final_flush bucket_bytes={bucket_bytes} bucket_len={len(bucket)}" + ) + if quantization and quantization.get("quant_method") == "fp8": + bucket = self._quantize_buffer_for_fp8( + bucket, + quantization_config=quantization, + target_device=self._fp8_quantization_target_device(backend), + phase_s=phase_s, + phase_prefix="direct_ep_fp8", + ) + t_backend = time.perf_counter() + backend.transfer_bucket( + bucket, + src_rank=self.rank, + flush_cache=flush_cache, + weight_version=weight_version, + ) + self._add_phase_time(phase_s, "direct_ep_backend_s", time.perf_counter() - t_backend) + num_buckets += 1 + logger.debug(f"Rank {self.rank}: [Direct-EP] {full_prefix} stage=after_final_flush") + + ctx["local_experts"] = None + logger.info( + f"Rank {self.rank}: [Direct-EP] {full_prefix} done " + f"total_bytes={total_bytes} total_params={total_params} num_buckets={num_buckets}" + ) + return total_bytes, total_params, num_buckets + + def _flush_pending_moe_bucket( + self, + backend, + flush_cache: bool = False, + weight_version: Optional[str] = None, + quantization: Optional[Dict[str, Any]] = None, + bucket_size_bytes: int = _DEFAULT_MOE_BUCKET_BYTES, + phase_s: Optional[Dict[str, float]] = None, + ) -> Tuple[int, int, int]: + """Ship the leftover MoE bucket accumulated across multiple + ``_direct_ep_transfer_experts(defer_final_flush=True)`` calls. + + Bytes/params already counted upstream by each ctx call (those + increment as tensors are appended to the bucket, regardless of + whether the bucket is shipped immediately or deferred), so we + only return the bucket count here. Returns (0, 0, 1) on a + non-empty bucket, (0, 0, 0) on empty. + """ + if self._pending_moe_cpu_workspace_records: + if not (quantization and quantization.get("quant_method") == "fp8"): + raise RuntimeError("FP8 CPU workspace records require FP8 quantization config") + bucket_bytes = self._pending_moe_bucket_bytes + nparams = len(self._pending_moe_cpu_workspace_records) + num_buckets = self._quantize_and_transfer_fp8_cpu_workspace_records( + backend, + self._pending_moe_cpu_workspace_records, + quantization_config=quantization, + bucket_size_bytes=bucket_size_bytes, + flush_cache=flush_cache, + weight_version=weight_version, + phase_s=phase_s, + phase_prefix="direct_ep_fp8", + ) + self._pending_moe_bucket = [] + self._pending_moe_bucket_bytes = 0 + self._reset_fp8_cpu_workspace_usage() + logger.info( + f"Rank {self.rank}: [WeightSync] Cross-layer MoE CPU workspace flush: " + f"{bucket_bytes / 1e6:.1f} MB source, {nparams} params, {num_buckets} transfer buckets" + ) + return 0, 0, num_buckets + + if not self._pending_moe_bucket: + self._pending_moe_bucket = [] + self._pending_moe_bucket_bytes = 0 + if flush_cache or weight_version is not None: + backend_config = getattr(getattr(backend, "config", None), "backend_config", None) + if backend_config is not None: + if weight_version is not None: + backend_config["weight_version"] = weight_version + backend_config["flush_cache"] = bool(flush_cache) + return 0, 0, 0 + bucket = self._pending_moe_bucket + bucket_bytes = self._pending_moe_bucket_bytes + nparams = len(bucket) + if quantization and quantization.get("quant_method") == "fp8": + bucket = self._quantize_buffer_for_fp8( + bucket, + quantization_config=quantization, + target_device=self._fp8_quantization_target_device(backend), + phase_s=phase_s, + phase_prefix="direct_ep_fp8", + ) + t_backend = time.perf_counter() + backend.transfer_bucket( + bucket, + src_rank=self.rank, + flush_cache=flush_cache, + weight_version=weight_version, + ) + self._add_phase_time(phase_s, "direct_ep_backend_s", time.perf_counter() - t_backend) + # Reset state for the next sync. + self._pending_moe_bucket = [] + self._pending_moe_bucket_bytes = 0 + logger.info( + f"Rank {self.rank}: [WeightSync] Cross-layer MoE flush: {bucket_bytes / 1e6:.1f} MB, {nparams} params" + ) + return 0, 0, 1 + + def _transfer_bucket_in_chunks( + self, + backend, + bucket: List[Tuple[str, torch.Tensor]], + *, + bucket_size_bytes: int, + flush_cache: bool, + weight_version: Optional[str], + phase_s: Optional[Dict[str, float]], + ) -> int: + num_buckets = 0 + chunk: List[Tuple[str, torch.Tensor]] = [] + chunk_bytes = 0 + + for name, tensor in bucket: + entry_bytes = tensor.numel() * tensor.element_size() + if chunk and chunk_bytes + entry_bytes > bucket_size_bytes: + t_backend = time.perf_counter() + backend.transfer_bucket( + chunk, + src_rank=self.rank, + flush_cache=False, + ) + self._add_phase_time(phase_s, "direct_ep_backend_s", time.perf_counter() - t_backend) + num_buckets += 1 + chunk = [] + chunk_bytes = 0 + + chunk.append((name, tensor)) + chunk_bytes += entry_bytes + + if chunk: + t_backend = time.perf_counter() + backend.transfer_bucket( + chunk, + src_rank=self.rank, + flush_cache=flush_cache, + weight_version=weight_version, + ) + self._add_phase_time(phase_s, "direct_ep_backend_s", time.perf_counter() - t_backend) + num_buckets += 1 + + return num_buckets + + def _quantize_and_transfer_fp8_cpu_workspace_records( + self, + backend, + records: List[Tuple[str, Tuple[Any, ...], int]], + *, + quantization_config: Dict[str, Any], + bucket_size_bytes: int, + flush_cache: bool, + weight_version: Optional[str], + phase_s: Optional[Dict[str, float]], + phase_prefix: str, + ) -> int: + stream_bytes = self._fp8_cpu_workspace_stream_bytes(bucket_size_bytes) + record_chunks = self._chunk_fp8_cpu_workspace_records(records, max_bytes=stream_bytes) + if not record_chunks: + return 0 + + if len(record_chunks) == 1 or not self._fp8_cpu_workspace_streaming_enabled(): + bucket = self._quantize_fp8_cpu_workspace_records( + records, + quantization_config=quantization_config, + phase_s=phase_s, + phase_prefix=phase_prefix, + ) + return self._transfer_bucket_in_chunks( + backend, + bucket, + bucket_size_bytes=bucket_size_bytes, + flush_cache=flush_cache, + weight_version=weight_version, + phase_s=phase_s, + ) + + logger.info( + "Rank %d: [WeightSync] Streaming FP8 CPU workspace flush in %d chunks (chunk cap %.1f MB)", + self.rank, + len(record_chunks), + stream_bytes / 1e6, + ) + + def transfer_task(bucket: List[Tuple[str, torch.Tensor]], is_final: bool) -> float: + t_backend = time.perf_counter() + backend.transfer_bucket( + bucket, + src_rank=self.rank, + flush_cache=(flush_cache if is_final else False), + weight_version=(weight_version if is_final else None), + ) + return time.perf_counter() - t_backend + + futures: List[Future[float]] = [] + with ThreadPoolExecutor(max_workers=1, thread_name_prefix=f"fp8-workspace-transfer-r{self.rank}") as executor: + for chunk_idx, record_chunk in enumerate(record_chunks): + bucket = self._quantize_fp8_cpu_workspace_records( + record_chunk, + quantization_config=quantization_config, + phase_s=phase_s, + phase_prefix=phase_prefix, + ) + futures.append(executor.submit(transfer_task, bucket, chunk_idx == len(record_chunks) - 1)) + + t_wait = time.perf_counter() + first_error: Optional[BaseException] = None + for future in futures: + try: + elapsed = future.result() + except BaseException as e: + if first_error is None: + first_error = e + continue + self._add_phase_time(phase_s, "direct_ep_backend_s", elapsed) + self._add_phase_time( + phase_s, + "direct_ep_fp8_workspace_stream_wait_s", + time.perf_counter() - t_wait, + ) + if first_error is not None: + for future in futures: + future.cancel() + raise first_error + + return len(record_chunks) + def _gather_and_broadcast_ep_moe_experts( self, backend, @@ -790,9 +2338,7 @@ def _gather_and_broadcast_ep_moe_experts( local_merged = [] for i in range(E_local): base_w = mod.dequantize_expert(proj_name, i, K, N) - a_idx = min(i, lora_A.shape[0] - 1) - b_idx = min(i, lora_B.shape[0] - 1) - delta = (lora_A[a_idx] @ lora_B[b_idx]) * mod.scaling + delta = self._compute_moe_lora_delta(mod, lora_A, lora_B, expert_idx=i) merged = base_w.to(torch.bfloat16) + delta.to(torch.bfloat16) local_merged.append(merged.t().contiguous()) # [N, K] del base_w, delta, merged @@ -835,6 +2381,7 @@ def _gather_and_broadcast_ep_moe_experts( bucket = self._quantize_buffer_for_fp8( bucket, quantization_config=quantization, + target_device=self._fp8_quantization_target_device(backend), ) b, p = self._broadcast_buffer( backend, @@ -855,6 +2402,7 @@ def _gather_and_broadcast_ep_moe_experts( bucket = self._quantize_buffer_for_fp8( bucket, quantization_config=quantization, + target_device=self._fp8_quantization_target_device(backend), ) b, p = self._broadcast_buffer( backend, @@ -933,9 +2481,7 @@ def _broadcast_moe_experts_bucketed( base_w = mod.dequantize_expert(proj_name, expert_idx, K, N) # [K, N] # Per-expert LoRA delta: A[i] @ B[i] * scaling → [K, N] - a_idx = min(expert_idx, lora_A.shape[0] - 1) - b_idx = min(expert_idx, lora_B.shape[0] - 1) - delta = (lora_A[a_idx] @ lora_B[b_idx]) * mod.scaling # [K, N] + delta = self._compute_moe_lora_delta(mod, lora_A, lora_B, expert_idx=expert_idx) # [K, N] merged = base_w.to(torch.bfloat16) + delta.to(torch.bfloat16) del base_w, delta @@ -955,6 +2501,7 @@ def _broadcast_moe_experts_bucketed( bucket = self._quantize_buffer_for_fp8( bucket, quantization_config=quantization, + target_device=self._fp8_quantization_target_device(backend), ) b, p = self._broadcast_buffer( backend, @@ -973,6 +2520,7 @@ def _broadcast_moe_experts_bucketed( bucket = self._quantize_buffer_for_fp8( bucket, quantization_config=quantization, + target_device=self._fp8_quantization_target_device(backend), ) b, p = self._broadcast_buffer( backend, @@ -1084,9 +2632,7 @@ def _compute_moe_experts_buffer( lora_B = lora_params[f"{hf_proj}_lora_B"] for expert_idx in range(E): base_w = mod.dequantize_expert(proj_name, expert_idx, K, N) - a_idx = min(expert_idx, lora_A.shape[0] - 1) - b_idx = min(expert_idx, lora_B.shape[0] - 1) - delta = (lora_A[a_idx] @ lora_B[b_idx]) * mod.scaling + delta = self._compute_moe_lora_delta(mod, lora_A, lora_B, expert_idx=expert_idx) merged = (base_w.to(torch.bfloat16) + delta.to(torch.bfloat16)).t().contiguous() del base_w, delta hf_name = f"{full_prefix}.{expert_idx}.{hf_proj}.weight" @@ -1104,10 +2650,716 @@ def _compute_moe_experts_buffer( # FP8 quantization # ======================================================================== + @staticmethod + def _fp8_quantization_target_device(backend) -> Optional[str]: + """Use CPU quantization for P2P, preserving NCCL's device-local path.""" + if backend.__class__.__name__ == "P2PTransportBackend": + return "cpu" + return None + + @staticmethod + def _fp8_quantization_execution_device() -> str: + return os.environ.get("XORL_P2P_FP8_QUANTIZE_DEVICE", "cpu").strip().lower() + + @staticmethod + def _fp8_cpu_workspace_enabled() -> bool: + return os.environ.get("XORL_P2P_FP8_CPU_WORKSPACE", "0") == "1" + + @staticmethod + def _fp8_cpu_workspace_pin_input() -> bool: + return os.environ.get("XORL_P2P_FP8_CPU_WORKSPACE_PINNED", "1") != "0" + + @staticmethod + def _fp8_cpu_workspace_min_capacity() -> int: + return _env_int("XORL_P2P_FP8_CPU_WORKSPACE_MIN_CAPACITY", 16) + + @staticmethod + def _fp8_cpu_workspace_streaming_enabled() -> bool: + return os.environ.get("XORL_P2P_FP8_CPU_WORKSPACE_STREAMING", "1") != "0" + + @staticmethod + def _p2p_requires_post_process_weights(quantization: Optional[Dict[str, Any]]) -> bool: + return bool(quantization and quantization.get("quant_method") == "fp8") + + @staticmethod + def _fp8_dtype_and_max(quantization_config: Dict[str, Any]) -> Tuple[torch.dtype, float]: + fmt = quantization_config.get("fmt", "e4m3") + if fmt == "e5m2": + fp8_dtype = torch.float8_e5m2 + else: + fp8_dtype = torch.float8_e4m3fn + return fp8_dtype, torch.finfo(fp8_dtype).max + + @staticmethod + def _fp8_block_size(quantization_config: Dict[str, Any]) -> Tuple[int, int]: + block_size_list = quantization_config.get("weight_block_size", [128, 128]) + block_size_row = block_size_list[0] + block_size_col = block_size_list[1] if len(block_size_list) > 1 else block_size_list[0] + return block_size_row, block_size_col + + def _reset_fp8_cpu_workspace_usage(self) -> None: + self._pending_moe_cpu_workspace_records = [] + for workspace in self._fp8_cpu_workspaces.values(): + workspace["used"] = 0 + + @staticmethod + def _add_phase_time(phase_s: Optional[Dict[str, float]], name: str, elapsed_s: float) -> None: + if phase_s is not None: + phase_s[name] = phase_s.get(name, 0.0) + elapsed_s + + @staticmethod + def _copy_tensor_to_cpu_for_fp8(tensor: torch.Tensor) -> torch.Tensor: + if tensor.device.type == "cpu": + return tensor.detach() + if ( + tensor.device.type == "cuda" + and torch.cuda.is_available() + and os.environ.get("XORL_P2P_FP8_PINNED_CPU_COPY", "1") != "0" + ): + cpu_tensor = torch.empty_like(tensor, device="cpu", pin_memory=True) + cpu_tensor.copy_(tensor.detach(), non_blocking=True) + torch.cuda.current_stream(tensor.device).synchronize() + return cpu_tensor + return tensor.detach().to("cpu") + + @staticmethod + def _should_quantize_fp8_weight( + name: str, + tensor: torch.Tensor, + modules_to_not_convert: List[str], + ) -> bool: + if not (name.endswith(".weight") and tensor.ndim == 2): + return False + if tensor.dtype in (torch.float8_e4m3fn, torch.float8_e5m2): + return False + + if modules_to_not_convert: + return not any( + name == prefix + ".weight" or name.startswith(prefix + ".") for prefix in modules_to_not_convert + ) + + if "_proj.weight" in name or name.endswith("fused_qkv_a_proj_with_mqa.weight"): + return True + + return ".linear_attn." in name and name.rsplit(".", 2)[-2] in { + "in_proj_qkv", + "in_proj_z", + "in_proj_b", + "in_proj_a", + } + + @staticmethod + def _can_group_fp8_tensor(first: torch.Tensor, tensor: torch.Tensor, group_len: int) -> bool: + if not (first.is_contiguous() and tensor.is_contiguous()): + return False + if first.shape != tensor.shape or first.dtype != tensor.dtype or first.device != tensor.device: + return False + if first.untyped_storage().data_ptr() != tensor.untyped_storage().data_ptr(): + return False + + rows, cols = first.shape + return tensor.storage_offset() == first.storage_offset() + group_len * rows * cols + + @staticmethod + def _can_quantize_fp8_stack_on_gpu( + stack: torch.Tensor, + *, + fp8_dtype: torch.dtype, + block_size_row: int, + block_size_col: int, + ) -> bool: + return ( + WeightSyncHandler._fp8_quantization_execution_device() in {"gpu", "cuda"} + and stack.device.type == "cuda" + and fp8_dtype == torch.float8_e4m3fn + and block_size_row == block_size_col + and stack.ndim == 3 + and stack.shape[1] % block_size_row == 0 + ) + + @staticmethod + def _quantize_fp8_stack_on_gpu_to_cpu( + stack: torch.Tensor, + *, + block_size: int, + phase_s: Optional[Dict[str, float]], + phase_prefix: str, + ) -> Tuple[torch.Tensor, torch.Tensor]: + from xorl.ops.quantize import block_fp8_quantize_gkn # noqa: PLC0415 + + if stack.ndim != 3: + raise ValueError(f"Expected a 3D FP8 quantization stack, got shape={tuple(stack.shape)}") + count, rows, cols = stack.shape + if rows % block_size != 0: + raise ValueError(f"GPU FP8 stack quantization requires rows divisible by {block_size}, got rows={rows}") + + t_quant = time.perf_counter() + work = stack.detach().contiguous() + flat = work.reshape(count * rows, cols) + quantized_flat, scale_flat = block_fp8_quantize_gkn(flat, block_size=block_size) + torch.cuda.current_stream(stack.device).synchronize() + WeightSyncHandler._add_phase_time(phase_s, f"{phase_prefix}_gpu_quant_s", time.perf_counter() - t_quant) + + scale_cols = (cols + block_size - 1) // block_size + quantized = quantized_flat.reshape(count, rows, cols) + scale_inv = scale_flat.reshape(count, rows // block_size, scale_cols) + + t_copy = time.perf_counter() + quantized_cpu = WeightSyncHandler._copy_tensor_to_cpu_for_fp8(quantized) + scale_cpu = WeightSyncHandler._copy_tensor_to_cpu_for_fp8(scale_inv) + WeightSyncHandler._add_phase_time( + phase_s, + f"{phase_prefix}_gpu_output_copy_s", + time.perf_counter() - t_copy, + ) + return quantized_cpu, scale_cpu + + @staticmethod + def _quantize_fp8_stack( + stack: torch.Tensor, + *, + fp8_dtype: torch.dtype, + fp8_max: float, + block_size_row: int, + block_size_col: int, + target_device: Optional[str], + phase_s: Optional[Dict[str, float]], + phase_prefix: str, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Quantize a [count, rows, cols] tensor stack and return FP8 weights + scales.""" + if stack.ndim != 3: + raise ValueError(f"Expected a 3D FP8 quantization stack, got shape={tuple(stack.shape)}") + + if ( + target_device is not None + and torch.device(target_device).type == "cpu" + and WeightSyncHandler._can_quantize_fp8_stack_on_gpu( + stack, + fp8_dtype=fp8_dtype, + block_size_row=block_size_row, + block_size_col=block_size_col, + ) + ): + return WeightSyncHandler._quantize_fp8_stack_on_gpu_to_cpu( + stack, + block_size=block_size_row, + phase_s=phase_s, + phase_prefix=phase_prefix, + ) + + work = stack.detach() + if target_device is not None: + t_copy = time.perf_counter() + if torch.device(target_device).type == "cpu": + work = WeightSyncHandler._copy_tensor_to_cpu_for_fp8(work) + else: + work = work.to(target_device) + WeightSyncHandler._add_phase_time( + phase_s, + f"{phase_prefix}_target_copy_s", + time.perf_counter() - t_copy, + ) + + t_float = time.perf_counter() + work = work.float() + WeightSyncHandler._add_phase_time(phase_s, f"{phase_prefix}_float_s", time.perf_counter() - t_float) + + count, rows, cols = work.shape + pad_rows = (block_size_row - rows % block_size_row) % block_size_row + pad_cols = (block_size_col - cols % block_size_col) % block_size_col + + if pad_rows > 0 or pad_cols > 0: + t_pad = time.perf_counter() + padded = torch.zeros( + count, + rows + pad_rows, + cols + pad_cols, + dtype=work.dtype, + device=work.device, + ) + padded[:, :rows, :cols] = work + WeightSyncHandler._add_phase_time(phase_s, f"{phase_prefix}_pad_s", time.perf_counter() - t_pad) + else: + padded = work + + nr = padded.shape[1] // block_size_row + nc = padded.shape[2] // block_size_col + blocks = padded.reshape(count, nr, block_size_row, nc, block_size_col).permute(0, 1, 3, 2, 4) + + t_reduce = time.perf_counter() + block_max = blocks.abs().reshape(count, nr, nc, -1).max(dim=-1).values + scale = block_max.clamp(min=1e-12) / fp8_max + scale_inv = scale.to(torch.float32) + WeightSyncHandler._add_phase_time(phase_s, f"{phase_prefix}_reduce_s", time.perf_counter() - t_reduce) + + t_cast = time.perf_counter() + scale_expanded = scale.unsqueeze(-1).unsqueeze(-1) + quantized_blocks = (blocks / scale_expanded).clamp(-fp8_max, fp8_max).to(fp8_dtype) + quantized = quantized_blocks.permute(0, 1, 3, 2, 4).reshape(count, padded.shape[1], padded.shape[2]) + WeightSyncHandler._add_phase_time(phase_s, f"{phase_prefix}_cast_s", time.perf_counter() - t_cast) + + if pad_rows > 0 or pad_cols > 0: + quantized = quantized[:, :rows, :cols].contiguous() + + return quantized, scale_inv + + @staticmethod + def _quantize_single_fp8_tensor( + tensor: torch.Tensor, + *, + fp8_dtype: torch.dtype, + fp8_max: float, + block_size_row: int, + block_size_col: int, + target_device: Optional[str], + phase_s: Optional[Dict[str, float]], + phase_prefix: str, + ) -> Tuple[torch.Tensor, torch.Tensor]: + quantized, scale_inv = WeightSyncHandler._quantize_fp8_stack( + tensor.unsqueeze(0), + fp8_dtype=fp8_dtype, + fp8_max=fp8_max, + block_size_row=block_size_row, + block_size_col=block_size_col, + target_device=target_device, + phase_s=phase_s, + phase_prefix=phase_prefix, + ) + return quantized[0].contiguous(), scale_inv[0].contiguous() + + @staticmethod + def _quantize_ep_expert_projection_for_fp8_cpu( + local_data: torch.Tensor, + *, + full_prefix: str, + proj_name: str, + ep_rank: int, + quantization_config: Dict[str, Any], + phase_s: Optional[Dict[str, float]], + ) -> Tuple[List[Tuple[str, torch.Tensor]], int]: + """CPU-quantize one EP-local MoE projection stack. + + ``local_data`` is stored in training layout [E, K, N]. SGLang locators + expect HF layout [N, K] per expert, plus a matching + ``weight_scale_inv`` tensor. + """ + entries, original_bytes = WeightSyncHandler._format_ep_expert_projection_for_fp8_cpu( + local_data, + full_prefix=full_prefix, + proj_name=proj_name, + ep_rank=ep_rank, + phase_s=phase_s, + ) + return ( + WeightSyncHandler._quantize_buffer_for_fp8( + entries, + quantization_config=quantization_config, + target_device=None, + phase_s=phase_s, + phase_prefix="direct_ep_fp8", + ), + original_bytes, + ) + + @staticmethod + def _format_ep_expert_projection_for_fp8_cpu( + local_data: torch.Tensor, + *, + full_prefix: str, + proj_name: str, + ep_rank: int, + phase_s: Optional[Dict[str, float]], + ) -> Tuple[List[Tuple[str, torch.Tensor]], int]: + """Copy one EP-local MoE projection to CPU HF layout without quantizing.""" + original_bytes = local_data.numel() * local_data.element_size() + t_copy = time.perf_counter() + cpu_data = WeightSyncHandler._copy_tensor_to_cpu_for_fp8(local_data) + WeightSyncHandler._add_phase_time(phase_s, "direct_ep_fp8_source_copy_s", time.perf_counter() - t_copy) + + t_transpose = time.perf_counter() + hf_stack = cpu_data.permute(0, 2, 1).contiguous() + WeightSyncHandler._add_phase_time(phase_s, "direct_ep_fp8_cpu_transpose_s", time.perf_counter() - t_transpose) + del cpu_data + + e_local = hf_stack.shape[0] + names = [f"{full_prefix}.{ep_rank * e_local + expert_idx}.{proj_name}.weight" for expert_idx in range(e_local)] + return [(name, hf_stack[idx]) for idx, name in enumerate(names)], original_bytes + + def _ensure_fp8_cpu_workspace( + self, + key: Tuple[Any, ...], + *, + required: int, + rows: int, + cols: int, + input_dtype: torch.dtype, + fp8_dtype: torch.dtype, + block_size_row: int, + block_size_col: int, + phase_s: Optional[Dict[str, float]], + ) -> Dict[str, Any]: + workspace = self._fp8_cpu_workspaces.get(key) + if workspace is not None and workspace["capacity"] >= required: + return workspace + + t_alloc = time.perf_counter() + old_workspace = workspace + old_used = int(old_workspace.get("used", 0)) if old_workspace is not None else 0 + old_capacity = int(old_workspace.get("capacity", 0)) if old_workspace is not None else 0 + min_capacity = self._fp8_cpu_workspace_min_capacity() + new_capacity = max(required, min_capacity, old_capacity * 2 if old_capacity else 0) + scale_rows = (rows + block_size_row - 1) // block_size_row + scale_cols = (cols + block_size_col - 1) // block_size_col + pin_input = self._fp8_cpu_workspace_pin_input() and torch.cuda.is_available() + + input_workspace = torch.empty( + (new_capacity, rows, cols), + dtype=input_dtype, + device="cpu", + pin_memory=pin_input, + ) + if old_workspace is not None and old_used: + input_workspace[:old_used].copy_(old_workspace["input"][:old_used]) + + workspace = { + "capacity": new_capacity, + "used": old_used, + "rows": rows, + "cols": cols, + "input_dtype": input_dtype, + "fp8_dtype": fp8_dtype, + "block_size_row": block_size_row, + "block_size_col": block_size_col, + "input": input_workspace, + "float": torch.empty((new_capacity, rows, cols), dtype=torch.float32, device="cpu"), + "abs": torch.empty((new_capacity, rows, cols), dtype=torch.float32, device="cpu"), + "quantized": torch.empty((new_capacity, rows, cols), dtype=fp8_dtype, device="cpu"), + "scale": torch.empty((new_capacity, scale_rows, scale_cols), dtype=torch.float32, device="cpu"), + } + self._fp8_cpu_workspaces[key] = workspace + self._add_phase_time(phase_s, "direct_ep_fp8_workspace_alloc_s", time.perf_counter() - t_alloc) + logger.info( + "Rank %d: [Direct-EP] FP8 CPU workspace key=%s capacity=%d rows=%d cols=%d pin_input=%s", + self.rank, + key, + new_capacity, + rows, + cols, + pin_input, + ) + return workspace + + def _stage_ep_expert_projection_for_fp8_cpu_workspace( + self, + local_data: torch.Tensor, + *, + full_prefix: str, + proj_name: str, + ep_rank: int, + quantization_config: Dict[str, Any], + phase_s: Optional[Dict[str, float]], + ) -> Tuple[List[Tuple[str, Tuple[Any, ...], int]], int]: + """Stage one EP-local projection into reusable CPU HF-layout storage.""" + fp8_dtype, _ = self._fp8_dtype_and_max(quantization_config) + block_size_row, block_size_col = self._fp8_block_size(quantization_config) + original_bytes = local_data.numel() * local_data.element_size() + e_local, cols, rows = local_data.shape + key = (rows, cols, local_data.dtype, fp8_dtype, block_size_row, block_size_col) + + workspace_used = int(self._fp8_cpu_workspaces.get(key, {}).get("used", 0)) + workspace = self._ensure_fp8_cpu_workspace( + key, + required=workspace_used + e_local, + rows=rows, + cols=cols, + input_dtype=local_data.dtype, + fp8_dtype=fp8_dtype, + block_size_row=block_size_row, + block_size_col=block_size_col, + phase_s=phase_s, + ) + start_idx = int(workspace_used) + end_idx = start_idx + e_local + + t_copy = time.perf_counter() + src = local_data.detach().permute(0, 2, 1) + dst = workspace["input"][start_idx:end_idx] + dst.copy_(src, non_blocking=(local_data.device.type == "cuda" and dst.is_pinned())) + if local_data.device.type == "cuda": + torch.cuda.current_stream(local_data.device).synchronize() + self._add_phase_time(phase_s, "direct_ep_fp8_workspace_copy_s", time.perf_counter() - t_copy) + + workspace["used"] = end_idx + records = [ + (f"{full_prefix}.{ep_rank * e_local + expert_idx}.{proj_name}.weight", key, start_idx + expert_idx) + for expert_idx in range(e_local) + ] + return records, original_bytes + + def _fp8_cpu_workspace_record_bytes(self, key: Tuple[Any, ...]) -> int: + workspace = self._fp8_cpu_workspaces[key] + rows = int(workspace["rows"]) + cols = int(workspace["cols"]) + block_size_row = int(workspace["block_size_row"]) + block_size_col = int(workspace["block_size_col"]) + fp8_dtype = workspace["fp8_dtype"] + scale_rows = (rows + block_size_row - 1) // block_size_row + scale_cols = (cols + block_size_col - 1) // block_size_col + weight_bytes = rows * cols * torch.empty((), dtype=fp8_dtype).element_size() + scale_bytes = scale_rows * scale_cols * torch.empty((), dtype=torch.float32).element_size() + return weight_bytes + scale_bytes + + def _fp8_cpu_workspace_records_bytes(self, records: List[Tuple[str, Tuple[Any, ...], int]]) -> int: + return sum(self._fp8_cpu_workspace_record_bytes(key) for _, key, _ in records) + + @staticmethod + def _fp8_cpu_workspace_stream_bytes(bucket_size_bytes: int) -> int: + max_stream_bytes = _env_int( + "XORL_P2P_FP8_CPU_WORKSPACE_STREAM_BYTES", + bucket_size_bytes, + ) + return min(max_stream_bytes, bucket_size_bytes) + + @staticmethod + def _fp8_cpu_workspace_pending_source_bytes(bucket_size_bytes: int) -> int: + return max(1, _env_int("XORL_P2P_FP8_CPU_WORKSPACE_PENDING_SOURCE_BYTES", bucket_size_bytes)) + + def _chunk_fp8_cpu_workspace_records( + self, + records: List[Tuple[str, Tuple[Any, ...], int]], + *, + max_bytes: int, + ) -> List[List[Tuple[str, Tuple[Any, ...], int]]]: + chunks: List[List[Tuple[str, Tuple[Any, ...], int]]] = [] + chunk: List[Tuple[str, Tuple[Any, ...], int]] = [] + chunk_bytes = 0 + + for record in records: + entry_bytes = self._fp8_cpu_workspace_record_bytes(record[1]) + if chunk and chunk_bytes + entry_bytes > max_bytes: + chunks.append(chunk) + chunk = [] + chunk_bytes = 0 + chunk.append(record) + chunk_bytes += entry_bytes + + if chunk: + chunks.append(chunk) + return chunks + + def _quantize_fp8_cpu_workspace_range( + self, + workspace: Dict[str, Any], + *, + start: int, + end: int, + fp8_dtype: torch.dtype, + fp8_max: float, + block_size_row: int, + block_size_col: int, + phase_s: Optional[Dict[str, float]], + phase_prefix: str, + ) -> None: + rows = int(workspace["rows"]) + cols = int(workspace["cols"]) + source = workspace["input"][start:end] + if rows % block_size_row != 0 or cols % block_size_col != 0: + quantized, scale = self._quantize_fp8_stack( + source, + fp8_dtype=fp8_dtype, + fp8_max=fp8_max, + block_size_row=block_size_row, + block_size_col=block_size_col, + target_device=None, + phase_s=phase_s, + phase_prefix=phase_prefix, + ) + workspace["quantized"][start:end].copy_(quantized) + workspace["scale"][start:end, : scale.shape[1], : scale.shape[2]].copy_(scale) + return + + count = end - start + work = workspace["float"][start:end] + abs_work = workspace["abs"][start:end] + scale = workspace["scale"][start:end] + quantized = workspace["quantized"][start:end] + nr = rows // block_size_row + nc = cols // block_size_col + + t_float = time.perf_counter() + work.copy_(source) + self._add_phase_time(phase_s, f"{phase_prefix}_float_s", time.perf_counter() - t_float) + + t_reduce = time.perf_counter() + torch.abs(work, out=abs_work) + blocks_abs = abs_work.reshape(count, nr, block_size_row, nc, block_size_col) + torch.amax(blocks_abs, dim=(2, 4), out=scale) + scale.clamp_(min=1e-12).div_(fp8_max) + self._add_phase_time(phase_s, f"{phase_prefix}_reduce_s", time.perf_counter() - t_reduce) + + t_cast = time.perf_counter() + blocks = work.reshape(count, nr, block_size_row, nc, block_size_col) + blocks.div_(scale.reshape(count, nr, 1, nc, 1)) + work.clamp_(min=-fp8_max, max=fp8_max) + quantized.copy_(work) + self._add_phase_time(phase_s, f"{phase_prefix}_cast_s", time.perf_counter() - t_cast) + + def _quantize_fp8_cpu_workspace_record_batch( + self, + records: List[Tuple[str, Tuple[Any, ...], int]], + *, + quantization_config: Dict[str, Any], + phase_s: Optional[Dict[str, float]], + phase_prefix: str, + ) -> None: + fp8_dtype, fp8_max = self._fp8_dtype_and_max(quantization_config) + block_size_row, block_size_col = self._fp8_block_size(quantization_config) + by_key: Dict[Tuple[Any, ...], List[int]] = {} + for _, key, index in records: + by_key.setdefault(key, []).append(index) + + for key, indices in by_key.items(): + workspace = self._fp8_cpu_workspaces[key] + used = int(workspace["used"]) + if workspace["fp8_dtype"] != fp8_dtype: + raise RuntimeError(f"FP8 workspace dtype mismatch: {workspace['fp8_dtype']} != {fp8_dtype}") + if int(workspace["block_size_row"]) != block_size_row or int(workspace["block_size_col"]) != block_size_col: + raise RuntimeError("FP8 workspace block-size mismatch") + + unique_indices = sorted(set(indices)) + if not unique_indices: + continue + if unique_indices[-1] >= used: + raise RuntimeError(f"FP8 workspace record index {unique_indices[-1]} exceeds used count {used}") + + range_start = unique_indices[0] + range_end = range_start + 1 + for index in unique_indices[1:]: + if index == range_end: + range_end += 1 + continue + self._quantize_fp8_cpu_workspace_range( + workspace, + start=range_start, + end=range_end, + fp8_dtype=fp8_dtype, + fp8_max=fp8_max, + block_size_row=block_size_row, + block_size_col=block_size_col, + phase_s=phase_s, + phase_prefix=phase_prefix, + ) + range_start = index + range_end = index + 1 + self._quantize_fp8_cpu_workspace_range( + workspace, + start=range_start, + end=range_end, + fp8_dtype=fp8_dtype, + fp8_max=fp8_max, + block_size_row=block_size_row, + block_size_col=block_size_col, + phase_s=phase_s, + phase_prefix=phase_prefix, + ) + + def _quantize_fp8_cpu_workspace_records( + self, + records: List[Tuple[str, Tuple[Any, ...], int]], + *, + quantization_config: Dict[str, Any], + phase_s: Optional[Dict[str, float]], + phase_prefix: str, + ) -> List[Tuple[str, torch.Tensor]]: + """Quantize staged workspace tensors while preserving record order.""" + self._quantize_fp8_cpu_workspace_record_batch( + records, + quantization_config=quantization_config, + phase_s=phase_s, + phase_prefix=phase_prefix, + ) + + result: List[Tuple[str, torch.Tensor]] = [] + for name, key, index in records: + workspace = self._fp8_cpu_workspaces[key] + result.append((name, workspace["quantized"][index])) + result.append((name.replace(".weight", ".weight_scale_inv"), workspace["scale"][index])) + return result + + @staticmethod + def _quantize_ep_expert_projection_for_fp8_gpu_to_cpu( + local_data: torch.Tensor, + *, + full_prefix: str, + proj_name: str, + ep_rank: int, + quantization_config: Dict[str, Any], + phase_s: Optional[Dict[str, float]], + ) -> Tuple[List[Tuple[str, torch.Tensor]], int]: + """GPU-quantize one EP-local MoE projection stack and return CPU tensors for P2P.""" + fmt = quantization_config.get("fmt", "e4m3") + if fmt != "e4m3": + raise ValueError("GPU FP8 quantization currently supports only e4m3") + + block_size_list = quantization_config.get("weight_block_size", [128, 128]) + block_size_row = block_size_list[0] + block_size_col = block_size_list[1] if len(block_size_list) > 1 else block_size_list[0] + if block_size_row != block_size_col: + raise ValueError("GPU FP8 quantization requires a square block size") + + original_bytes = local_data.numel() * local_data.element_size() + + t_layout = time.perf_counter() + hf_stack = local_data.detach().permute(0, 2, 1).contiguous() + torch.cuda.current_stream(local_data.device).synchronize() + WeightSyncHandler._add_phase_time(phase_s, "direct_ep_fp8_gpu_layout_s", time.perf_counter() - t_layout) + + e_local = hf_stack.shape[0] + modules_to_not_convert = quantization_config.get("modules_to_not_convert", []) + names = [f"{full_prefix}.{ep_rank * e_local + idx}.{proj_name}.weight" for idx in range(e_local)] + if not all( + WeightSyncHandler._should_quantize_fp8_weight(name, hf_stack[idx], modules_to_not_convert) + for idx, name in enumerate(names) + ): + t_copy = time.perf_counter() + hf_stack_cpu = WeightSyncHandler._copy_tensor_to_cpu_for_fp8(hf_stack) + WeightSyncHandler._add_phase_time( + phase_s, + "direct_ep_fp8_gpu_output_copy_s", + time.perf_counter() - t_copy, + ) + entries = [(name, hf_stack_cpu[idx]) for idx, name in enumerate(names)] + return ( + WeightSyncHandler._quantize_buffer_for_fp8( + entries, + quantization_config=quantization_config, + target_device=None, + phase_s=phase_s, + phase_prefix="direct_ep_fp8", + ), + original_bytes, + ) + + quantized_stack, scale_stack = WeightSyncHandler._quantize_fp8_stack_on_gpu_to_cpu( + hf_stack, + block_size=block_size_row, + phase_s=phase_s, + phase_prefix="direct_ep_fp8", + ) + + result: List[Tuple[str, torch.Tensor]] = [] + for idx, name in enumerate(names): + result.append((name, quantized_stack[idx])) + result.append((name.replace(".weight", ".weight_scale_inv"), scale_stack[idx])) + return result, original_bytes + @staticmethod def _quantize_buffer_for_fp8( buffer: List[Tuple[str, torch.Tensor]], quantization_config: Optional[Dict[str, Any]] = None, + target_device: Optional[str] = None, + phase_s: Optional[Dict[str, float]] = None, + phase_prefix: str = "fp8", ) -> List[Tuple[str, torch.Tensor]]: """Quantize bf16 weight tensors to FP8 with block-wise scales. @@ -1121,91 +3373,79 @@ def _quantize_buffer_for_fp8( - modules_to_not_convert: list of module name prefixes to skip quantization Non-weight params (e.g. layernorm, embedding) are passed through as-is. + When target_device="cpu", quantized tensors are returned on CPU so the + P2P backend can stage them through the registered CPU pool and transfer + fewer bytes to the receiver. """ if quantization_config is None: quantization_config = {} # FP8 format: e4m3 (default, higher precision) or e5m2 (wider range) - fmt = quantization_config.get("fmt", "e4m3") - if fmt == "e5m2": - fp8_dtype = torch.float8_e5m2 - else: - fp8_dtype = torch.float8_e4m3fn - fp8_max = torch.finfo(fp8_dtype).max - - block_size_list = quantization_config.get("weight_block_size", [128, 128]) - block_size_row = block_size_list[0] - block_size_col = block_size_list[1] if len(block_size_list) > 1 else block_size_list[0] + fp8_dtype, fp8_max = WeightSyncHandler._fp8_dtype_and_max(quantization_config) + block_size_row, block_size_col = WeightSyncHandler._fp8_block_size(quantization_config) modules_to_not_convert = quantization_config.get("modules_to_not_convert", []) + target_is_cpu = target_device is not None and torch.device(target_device).type == "cpu" + result = [] - for name, tensor in buffer: - # Must be a 2D weight tensor to be quantized - if not (name.endswith(".weight") and tensor.ndim == 2): + i = 0 + while i < len(buffer): + name, tensor = buffer[i] + if not WeightSyncHandler._should_quantize_fp8_weight(name, tensor, modules_to_not_convert): result.append((name, tensor)) + i += 1 continue - # Check modules_to_not_convert: match if the param name starts with - # any entry (prefix match). E.g. "lm_head" matches "lm_head.weight", - # "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 + group_end = i + 1 + if tensor.device.type == "cpu" or (target_is_cpu and tensor.device.type != "cpu"): + while group_end < len(buffer): + next_name, next_tensor = buffer[group_end] + if not WeightSyncHandler._should_quantize_fp8_weight( + next_name, + next_tensor, + modules_to_not_convert, + ): + break + if not WeightSyncHandler._can_group_fp8_tensor(tensor, next_tensor, group_end - i): + break + group_end += 1 + + if group_end > i + 1: + rows, cols = tensor.shape + stack = torch.as_strided( + tensor, + size=(group_end - i, rows, cols), + stride=(rows * cols, cols, 1), ) - if skip: - result.append((name, tensor)) - continue - else: - # Default skip logic when no explicit list: only quantize _proj weights - if "_proj.weight" not in name: - result.append((name, tensor)) - continue - - rows, cols = tensor.shape - # Pad to block_size alignment if needed - pad_rows = (block_size_row - rows % block_size_row) % block_size_row - pad_cols = (block_size_col - cols % block_size_col) % block_size_col - - if pad_rows > 0 or pad_cols > 0: - padded = torch.zeros( - rows + pad_rows, - cols + pad_cols, - dtype=tensor.dtype, - device=tensor.device, + quantized_stack, scale_stack = WeightSyncHandler._quantize_fp8_stack( + stack, + fp8_dtype=fp8_dtype, + fp8_max=fp8_max, + block_size_row=block_size_row, + block_size_col=block_size_col, + target_device=target_device, + phase_s=phase_s, + phase_prefix=phase_prefix, ) - padded[:rows, :cols] = tensor - else: - padded = tensor - - # Reshape into blocks: [nr, block_size_row, nc, block_size_col] - nr = padded.shape[0] // block_size_row - nc = padded.shape[1] // block_size_col - blocks = padded.reshape(nr, block_size_row, nc, block_size_col).permute(0, 2, 1, 3) - # blocks shape: [nr, nc, block_size_row, block_size_col] - - # Compute per-block scale: max(abs(block)) / fp8_max - block_max = blocks.abs().reshape(nr, nc, -1).max(dim=-1).values # [nr, nc] - scale = block_max.clamp(min=1e-12) / fp8_max # [nr, nc] - scale_inv = scale.to(torch.float32) - - # Quantize: divide by scale, clamp, cast to fp8 - # Expand scale for broadcasting: [nr, nc, 1, 1] - scale_expanded = scale.unsqueeze(-1).unsqueeze(-1) # [nr, nc, 1, 1] - quantized_blocks = (blocks.float() / scale_expanded).clamp(-fp8_max, fp8_max) - 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]) - - # Remove padding - if pad_rows > 0 or pad_cols > 0: - quantized = quantized[:rows, :cols].contiguous() - - # scale_inv name: replace .weight with .weight_scale_inv - scale_name = name.replace(".weight", ".weight_scale_inv") + for group_idx, (group_name, _) in enumerate(buffer[i:group_end]): + result.append((group_name, quantized_stack[group_idx])) + result.append((group_name.replace(".weight", ".weight_scale_inv"), scale_stack[group_idx])) + i = group_end + continue + quantized, scale_inv = WeightSyncHandler._quantize_single_fp8_tensor( + tensor, + fp8_dtype=fp8_dtype, + fp8_max=fp8_max, + block_size_row=block_size_row, + block_size_col=block_size_col, + target_device=target_device, + phase_s=phase_s, + phase_prefix=phase_prefix, + ) result.append((name, quantized)) - result.append((scale_name, scale_inv)) + result.append((name.replace(".weight", ".weight_scale_inv"), scale_inv)) + i += 1 return result @@ -1213,12 +3453,47 @@ def _quantize_buffer_for_fp8( # Helpers # ======================================================================== + @staticmethod + def _would_exceed_bucket_cap( + current_bytes: int, + next_entry_bytes: int, + bucket_size_bytes: int, + ) -> bool: + return current_bytes > 0 and current_bytes + next_entry_bytes > bucket_size_bytes + + @staticmethod + def _resolve_module_path(root: Any, path: str) -> Optional[Any]: + current = root + for part in path.split("."): + if not part: + continue + if part.isdigit() and hasattr(current, "__getitem__"): + try: + current = current[int(part)] + continue + except Exception: + return None + try: + current = getattr(current, part) + except AttributeError: + return None + return current + + @staticmethod + def _module_paths_share_parameter(root: Any, tied_name: str, source_name: str) -> bool: + tied = WeightSyncHandler._resolve_module_path(root, tied_name) + source = WeightSyncHandler._resolve_module_path(root, source_name) + return tied is not None and tied is source + @staticmethod def _extract_params_for_sync( fsdp_mod, mod_name: str, DTensor, skip_moe_prefixes: Optional[set] = None, + emit_tied_weight_duplicates: bool = True, + tied_weight_aliases: Optional[Dict[str, str]] = None, + include_param: Optional[Callable[[str], bool]] = None, ) -> List[Tuple[str, torch.Tensor]]: """ Extract parameters from an unsharded FSDP module for sync. @@ -1284,6 +3559,17 @@ def _extract_params_for_sync( continue full_name = f"{mod_name}.{pname}" if mod_name != "(root)" else pname + if include_param is not None and not include_param(full_name): + continue + + if not emit_tied_weight_duplicates and tied_weight_aliases is not None: + source_name = tied_weight_aliases.get(full_name) + if source_name is not None and source_name != full_name: + logger.info( + f"Rank 0: [WeightSync] Tied weight: deferring {full_name} " + f"to receiver-side copy from {source_name}" + ) + continue # Check if this is a base weight with LoRA to merge # Case 1: LoraLinear — pname like "self_attn.q_proj.weight" @@ -1335,7 +3621,8 @@ def _extract_params_for_sync( pass if _is_moe_experts: - buffer.append((full_name, param.data.to(dtype=torch.bfloat16).clone())) + cloned_moe = param.data.to(dtype=torch.bfloat16).clone() + buffer.append((full_name, cloned_moe)) else: cloned = param.data.to(dtype=torch.bfloat16).clone() buffer.append((full_name, cloned)) @@ -1352,10 +3639,27 @@ def _extract_params_for_sync( if full_tied not in buffer_names and full_source in buffer_names: for buf_name, buf_tensor in buffer: if buf_name == full_source: - logger.info( - f"Rank 0: [WeightSync] Tied weight: emitting {full_tied} (clone of {full_source})" - ) - buffer.append((full_tied, buf_tensor.clone())) + if emit_tied_weight_duplicates: + logger.info( + f"Rank 0: [WeightSync] Tied weight: emitting {full_tied} (clone of {full_source})" + ) + buffer.append((full_tied, buf_tensor.clone())) + elif WeightSyncHandler._module_paths_share_parameter( + fsdp_mod, + tied_name, + source_name, + ): + if tied_weight_aliases is not None: + tied_weight_aliases[full_tied] = full_source + logger.info( + f"Rank 0: [WeightSync] Tied weight: deferring {full_tied} " + f"to receiver-side copy from {full_source}" + ) + else: + logger.info( + f"Rank 0: [WeightSync] Tied weight: not deferring {full_tied}; " + f"{tied_name} and {source_name} are not the same Parameter" + ) break return buffer @@ -1371,6 +3675,8 @@ def _unfuse_for_inference( - 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 + - DeepseekV3 / Kimi-K2.5 MLA: q_a_proj + kv_a_proj_with_mqa → + fused_qkv_a_proj_with_mqa to match SGLang's inference module - Qwen3.5 linear attention: remap split GatedDeltaNet params back to HF fused names (q_proj/k_proj/v_proj → in_proj_qkv, etc.) """ @@ -1415,11 +3721,54 @@ def _unfuse_for_inference( else: result.append((name, tensor)) + result = WeightSyncHandler._remap_deepseek_mla_params_for_inference(result, config) + if has_linear_attention_layers(config): result = remap_linear_attention_params_for_inference(result) return result + @staticmethod + def _remap_deepseek_mla_params_for_inference( + buffer: List[Tuple[str, torch.Tensor]], + config, + ) -> List[Tuple[str, torch.Tensor]]: + """Fuse DeepseekV3/Kimi-K2.5 MLA A projections for SGLang receivers.""" + if getattr(config, "q_lora_rank", None) is None: + return buffer + + result: List[Tuple[str, torch.Tensor]] = [] + pending: Dict[Tuple[str, str], Dict[str, torch.Tensor]] = {} + + for name, tensor in buffer: + if ".self_attn.q_a_proj." in name: + prefix, suffix = name.rsplit(".q_a_proj.", 1) + pending.setdefault((prefix, suffix), {})["q_a_proj"] = tensor + continue + if ".self_attn.kv_a_proj_with_mqa." in name: + prefix, suffix = name.rsplit(".kv_a_proj_with_mqa.", 1) + pending.setdefault((prefix, suffix), {})["kv_a_proj_with_mqa"] = tensor + continue + result.append((name, tensor)) + + for (prefix, suffix), parts in pending.items(): + q_a = parts.get("q_a_proj") + kv_a = parts.get("kv_a_proj_with_mqa") + if suffix == "weight" and q_a is not None and kv_a is not None: + result.append( + ( + f"{prefix}.fused_qkv_a_proj_with_mqa.{suffix}", + torch.cat([q_a, kv_a], dim=0).contiguous(), + ) + ) + continue + if q_a is not None: + result.append((f"{prefix}.q_a_proj.{suffix}", q_a)) + if kv_a is not None: + result.append((f"{prefix}.kv_a_proj_with_mqa.{suffix}", kv_a)) + + return result + def _broadcast_buffer( self, backend, @@ -1442,13 +3791,49 @@ def _broadcast_buffer( bucket_bytes = sum(t.numel() * t.element_size() for _, t in buffer) logger.info(f"Rank {self.rank}: [WeightSync] Broadcasting {len(buffer)} params, {bucket_bytes / 1e6:.1f} MB") - backend.transfer_bucket( - buffer, - flush_cache=flush_cache, - weight_version=weight_version, + dense_bucket_bytes = _env_int("XORL_WEIGHT_SYNC_DENSE_BUCKET_BYTES", 0, minimum=0) + if dense_bucket_bytes <= 0 or bucket_bytes <= dense_bucket_bytes: + backend.transfer_bucket( + buffer, + flush_cache=flush_cache, + weight_version=weight_version, + ) + return bucket_bytes, len(buffer) + + chunks = self._chunk_buffer_by_bytes(buffer, dense_bucket_bytes) + logger.info( + f"Rank {self.rank}: [WeightSync] Split dense buffer into {len(chunks)} transfer buckets " + f"(target={dense_bucket_bytes / 1e6:.1f} MB)" ) + for idx, chunk in enumerate(chunks): + is_last = idx == len(chunks) - 1 + backend.transfer_bucket( + chunk, + flush_cache=flush_cache and is_last, + weight_version=weight_version if is_last else None, + ) return bucket_bytes, len(buffer) + @staticmethod + def _chunk_buffer_by_bytes( + buffer: List[Tuple[str, torch.Tensor]], + bucket_size_bytes: int, + ) -> List[List[Tuple[str, torch.Tensor]]]: + chunks: List[List[Tuple[str, torch.Tensor]]] = [] + chunk: List[Tuple[str, torch.Tensor]] = [] + chunk_bytes = 0 + for name, tensor in buffer: + tensor_bytes = tensor.numel() * tensor.element_size() + if chunk and chunk_bytes + tensor_bytes > bucket_size_bytes: + chunks.append(chunk) + chunk = [] + chunk_bytes = 0 + chunk.append((name, tensor)) + chunk_bytes += tensor_bytes + if chunk: + chunks.append(chunk) + return chunks + @staticmethod def _get_fsdp_modules(model) -> Tuple[Optional[Any], List[Tuple[str, Any]]]: """ diff --git a/src/xorl/trainers/model_builder.py b/src/xorl/trainers/model_builder.py index 39644534..bcdff8a1 100644 --- a/src/xorl/trainers/model_builder.py +++ b/src/xorl/trainers/model_builder.py @@ -11,12 +11,17 @@ import torch import torch.nn as nn +from xorl.distributed.parallel_state import get_parallel_state 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.lora.utils import inject_lora_into_model, inject_lora_into_model_with_moe from xorl.models import build_foundation_model from xorl.models.checkpoint_handlers.buffers import get_prequantized_exclude_modules from xorl.models.layers.rope import set_rope_native +from xorl.models.transformers.deepseek_v3.support import ( + freeze_deepseek_v3_router_parameters, + validate_deepseek_v3_training_mode, +) from xorl.qlora import ( detect_prequantized_block_fp8, detect_prequantized_nvfp4, @@ -106,6 +111,7 @@ def build_training_model( moe_implementation: Optional[str] = None, ep_dispatch: str = "alltoall", train_router: bool = False, + record_routing_weights: bool = True, deepep_buffer_size_gb: float = 2.0, deepep_num_sms: int = 20, deepep_async_combine: bool = False, @@ -130,8 +136,9 @@ def build_training_model( basic_modules: Optional[List[str]] = None, enable_reentrant: bool = False, enable_forward_prefetch: bool = True, - load_weights_mode: str = "broadcast", + load_weights_mode: str = "grouped", reshard_after_forward: Optional[bool] = None, + moe_grad_reduce_mode: str = "reduce_scatter", pp_schedule: Optional[str] = None, # --- Training flags --- freeze_router: bool = False, @@ -142,6 +149,7 @@ def build_training_model( activation_native: bool = False, rope_native: bool = False, attention_cast_bf16: bool = False, + flash_attention_deterministic: bool = False, ) -> TrainingModelResult: """Build, inject LoRA/QLoRA, and parallelize a training model. @@ -173,6 +181,7 @@ def build_training_model( moe_implementation=moe_implementation, ep_dispatch=ep_dispatch, train_router=train_router, + record_routing_weights=record_routing_weights, deepep_buffer_size_gb=deepep_buffer_size_gb, deepep_num_sms=deepep_num_sms, deepep_async_combine=deepep_async_combine, @@ -182,6 +191,7 @@ def build_training_model( activation_native=activation_native, rope_native=rope_native, attention_cast_bf16=attention_cast_bf16, + flash_attention_deterministic=flash_attention_deterministic, init_device=init_device, ) @@ -192,6 +202,12 @@ def build_training_model( if activation_native: logger.info_rank0("Using native SiLU activation (fused Triton kernel disabled)") model_config = model.config + validate_deepseek_v3_training_mode( + model_config, + enable_qlora=enable_qlora, + freeze_router=freeze_router, + merge_qkv=merge_qkv, + ) helper.print_device_mem_info("VRAM usage after building model") # ------------------------------------------------------------------ @@ -250,6 +266,7 @@ def build_training_model( # 6. Parallelize (FSDP2 / PP) # ------------------------------------------------------------------ _basic_modules = list(model._no_split_modules) + (basic_modules or []) + effective_pp_schedule = pp_schedule if get_parallel_state().pp_enabled else None build_result = _parallelize( model, init_device=init_device, @@ -262,8 +279,9 @@ def build_training_model( enable_reentrant=enable_reentrant, enable_forward_prefetch=enable_forward_prefetch, load_weights_mode=load_weights_mode, - pp_schedule=pp_schedule, + pp_schedule=effective_pp_schedule, reshard_after_forward=reshard_after_forward, + moe_grad_reduce_mode=moe_grad_reduce_mode, skip_param_upcast=should_skip_generic_param_upcast( enable_lora=enable_lora, enable_qlora=enable_qlora, @@ -309,11 +327,12 @@ def build_training_model( # Optionally freeze MoE router if freeze_router: - router_frozen_count = 0 - for name, param in model.named_parameters(): - if ".gate.weight" in name: - param.requires_grad = False - router_frozen_count += 1 + router_frozen_count = freeze_deepseek_v3_router_parameters(model) + if router_frozen_count == 0: + for name, param in model.named_parameters(): + if ".gate.weight" in name: + param.requires_grad = False + router_frozen_count += 1 if router_frozen_count > 0: logger.info_rank0(f"Froze {router_frozen_count} MoE router (gate) parameters") @@ -438,8 +457,6 @@ def _inject_lora( is_moe_model = getattr(model.config, "num_experts", 0) > 0 if is_moe_model and moe_hybrid_shared_lora: - from xorl.lora.utils import inject_lora_into_model_with_moe - logger.info_rank0(f"MoE-aware LoRA injection (hybrid_shared={moe_hybrid_shared_lora})") inject_lora_into_model_with_moe( model, @@ -462,7 +479,7 @@ def _inject_lora( def _deferred_qlora_quantize( model: nn.Module, weights_path: str, - load_weights_mode: str = "broadcast", + load_weights_mode: str = "grouped", ) -> None: """After FSDP loads weights, quantize/load weights into QLoRA modules. diff --git a/src/xorl/trainers/trainer.py b/src/xorl/trainers/trainer.py index dc58df5c..05d4579a 100644 --- a/src/xorl/trainers/trainer.py +++ b/src/xorl/trainers/trainer.py @@ -38,6 +38,10 @@ 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.deepseek_v3.support import ( + freeze_deepseek_v3_router_parameters, + validate_deepseek_v3_training_mode, +) from xorl.optim import build_lr_scheduler, build_optimizer from xorl.qlora import ( detect_prequantized_block_fp8, @@ -55,12 +59,16 @@ ) from xorl.trainers.training_utils import ( clip_gradients, + count_active_microbatches, count_valid_tokens, forward_backward_pp, + get_distsign_grad_scale_factor, + get_effective_grad_clip_value, + make_pp_loss_fn, maybe_merge_lora, negotiate_pp_seq_len, pad_micro_batches_for_pp, - pp_loss_fn, + scale_model_gradients, sync_sp_gradients, ) from xorl.utils import helper @@ -70,6 +78,13 @@ logger = helper.create_logger(__name__) _trainer_cpu_group: Optional[dist.ProcessGroup] = None +_HOST_INVENTORY_MAX_WORLD_SIZE = int(os.environ.get("XORL_HOST_INVENTORY_MAX_WORLD_SIZE", "64")) +_HOST_INVENTORY_DISABLED = os.environ.get("XORL_DISABLE_HOST_INVENTORY", "").strip().lower() in { + "1", + "true", + "yes", + "on", +} def _get_trainer_cpu_group() -> Optional[dist.ProcessGroup]: @@ -82,6 +97,12 @@ def _get_trainer_cpu_group() -> Optional[dist.ProcessGroup]: return _trainer_cpu_group +def _should_collect_host_inventory(world_size: int) -> bool: + if _HOST_INVENTORY_DISABLED: + return False + return world_size <= _HOST_INVENTORY_MAX_WORLD_SIZE + + # --------------------------------------------------------------------------- # TrainState — checkpointable training state # --------------------------------------------------------------------------- @@ -161,7 +182,7 @@ def _timed(phase_name, fn): def _bootstrap(self) -> None: """Initialize distributed, device, seed, parallel state.""" args = self.args - from datetime import timedelta + from datetime import timedelta # noqa: PLC0415 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}") @@ -231,6 +252,16 @@ def _bootstrap(self) -> None: # Routing replay is only needed with EP when MoE forward is recomputed self._use_routing_replay = self.ps.ep_size > 1 and args.train.moe_recomputed + # Loss-function kwargs forwarded to causallm_loss_function each step. + self._causallm_loss_params: Dict[str, Any] = {"ce_mode": args.train.ce_mode} + if args.train.softmax_auxiliary_loss: + if args.train.pipeline_parallel_size > 1: + raise NotImplementedError( + "softmax_auxiliary_loss (Z-loss) is not yet supported with pipeline parallelism. " + "PP uses a separate compiled CE loss path (pp_loss_fn) that does not compute logsumexp." + ) + self._causallm_loss_params["z_loss_coef"] = args.train.auxiliary_loss_multiplier + 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: @@ -267,6 +298,8 @@ def _collect_host_inventory(self) -> List[Dict[str, Any]]: } if not dist.is_available() or not dist.is_initialized() or self.args.train.world_size <= 1: return [payload] + if not _should_collect_host_inventory(self.args.train.world_size): + 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] @@ -277,6 +310,22 @@ def _log_host_inventory(self) -> None: if self.args.train.global_rank != 0: return + if self.args.train.world_size > 1 and not _should_collect_host_inventory(self.args.train.world_size): + logger.info_rank0( + "Skipping host inventory gather for world_size=%s; " + "set XORL_HOST_INVENTORY_MAX_WORLD_SIZE or XORL_DISABLE_HOST_INVENTORY to override.", + self.args.train.world_size, + ) + self._startup_metrics.update( + { + "startup/master_addr": os.environ.get("MASTER_ADDR"), + "startup/master_port": os.environ.get("MASTER_PORT"), + "startup/host_inventory_skipped": True, + "startup/host_inventory_world_size": self.args.train.world_size, + } + ) + 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( @@ -392,13 +441,24 @@ def _build_model(self) -> None: moe_implementation=args.model.moe_implementation, ep_dispatch=args.model.ep_dispatch, train_router=args.model.train_router, + record_routing_weights=args.model.record_routing_weights, 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, + flash_attention_deterministic=args.model.flash_attention_deterministic, init_device=args.train.init_device, ) self.model_config = self.model.config + # Normalize _no_split_modules to list — some HF models (e.g. GPT-OSS) define it as a set + if isinstance(getattr(self.model, "_no_split_modules", None), set): + self.model._no_split_modules = list(self.model._no_split_modules) + validate_deepseek_v3_training_mode( + self.model_config, + enable_qlora=args.lora.enable_qlora, + freeze_router=args.model.freeze_router, + merge_qkv=args.model.merge_qkv, + ) helper.print_device_mem_info("VRAM usage after building model") # Unfuse QKV for tensor parallelism @@ -480,7 +540,7 @@ def _inject_lora(self) -> None: 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 + from xorl.lora.utils import inject_lora_into_model_with_moe # noqa: PLC0415 logger.info_rank0(f"MoE-aware LoRA injection (hybrid_shared={args.lora.moe_hybrid_shared_lora})") inject_lora_into_model_with_moe( @@ -491,7 +551,7 @@ def _inject_lora(self) -> None: moe_hybrid_shared_lora=args.lora.moe_hybrid_shared_lora, ) else: - from xorl.lora.utils import inject_lora_into_model + from xorl.lora.utils import inject_lora_into_model # noqa: PLC0415 inject_lora_into_model( self.model, @@ -523,6 +583,7 @@ def _parallelize(self) -> None: load_weights_mode=args.train.load_weights_mode, pp_schedule=args.train.pipeline_parallel_schedule if args.train.pipeline_parallel_size > 1 else None, reshard_after_forward=args.train.reshard_after_forward, + moe_grad_reduce_mode=args.train.moe_grad_reduce_mode, skip_param_upcast=should_skip_generic_param_upcast( enable_lora=args.lora.enable_lora, enable_qlora=args.lora.enable_qlora, @@ -557,6 +618,16 @@ def _parallelize(self) -> None: if "lora_A" not in name and "lora_B" not in name: param.requires_grad = False + if args.model.freeze_router: + frozen = freeze_deepseek_v3_router_parameters(self.model) + if frozen == 0: + for name, param in self.model.named_parameters(): + if ".gate.weight" in name: + param.requires_grad = False + frozen += 1 + if frozen > 0: + logger.info_rank0(f"Froze {frozen} MoE router (gate) parameters") + def _deferred_qlora_quantize(self) -> None: """After FSDP loads weights, quantize them into uint8 buffers.""" args = self.args @@ -591,6 +662,7 @@ def _deferred_qlora_quantize(self) -> None: def _build_optimizer(self) -> None: """Build optimizer and LR scheduler.""" args = self.args + self._use_distsignsgd = args.train.optimizer == "distsignsgd" self.optimizer = build_optimizer( self.model, lr=args.train.lr, @@ -599,6 +671,7 @@ def _build_optimizer(self) -> None: optimizer_type=args.train.optimizer, optimizer_dtype=args.train.optimizer_dtype, optimizer_kwargs=args.train.optimizer_kwargs, + cautious_weight_decay=args.train.cautious_weight_decay, ) if self._optimizer_pre_hook_fn is not None: hook = self._optimizer_pre_hook_fn(self.model, self.model_config, args.train.data_parallel_mode) @@ -658,7 +731,27 @@ def _resume_checkpoint(self) -> None: if not args.train.load_checkpoint_path: return - state = {"model": self.model, "optimizer": self.optimizer, "extra_state": {}} + # When load_weights_mode=skip, parameters are meta-device DTensors after FSDP wrapping. + # Use to_empty() to materialize them to real CUDA tensors while preserving the DTensor + # wrapper (unlike manual setattr which would break FSDP2's internal state). + if args.train.load_weights_mode == "skip": + logger.info_rank0("Materializing meta parameters to CUDA via to_empty()...") + self.model.to_empty(device=f"cuda:{args.train.local_rank}") + logger.info_rank0("Meta parameters materialized.") + + state = {"model": self.model, "extra_state": {}} + # Only include optimizer if the checkpoint has optimizer state (i.e., resuming training). + # Model-only DCP checkpoints (from convert_checkpoint.py) won't have optimizer state. + ckpt_has_optimizer = os.path.exists(os.path.join(args.train.load_checkpoint_path, ".metadata")) + if ckpt_has_optimizer: + import torch.distributed.checkpoint as dcp_meta # noqa: PLC0415 + + try: + metadata = dcp_meta.FileSystemReader(args.train.load_checkpoint_path).read_metadata() + if any(k.startswith("optimizer") for k in metadata.state_dict_metadata.keys()): + state["optimizer"] = self.optimizer + except Exception: + pass self.Checkpointer.load(args.train.load_checkpoint_path, state) extra = state.get("extra_state", {}) @@ -702,7 +795,7 @@ def _get_pp_schedule(self, seq_len: int): self._pp_schedule_cache[seq_len] = build_pipeline_schedule( stages=[stage], n_microbatches=self.args.train.gradient_accumulation_steps, - loss_fn=pp_loss_fn, + loss_fn=make_pp_loss_fn(self.args.train.ce_mode), schedule_name=self.args.train.pipeline_parallel_schedule, ) return self._pp_schedule_cache[seq_len] @@ -829,6 +922,10 @@ def train_step(self, micro_batches: List[Dict[str, Any]]) -> Tuple[float, float] (total_loss, grad_norm) — all-reduced across DP for logging. """ global_valid_tokens = self._count_valid_tokens(micro_batches) + if self._use_distsignsgd: + active_microbatches, active_voter_total = self._count_active_microbatches(micro_batches) + else: + active_microbatches, active_voter_total = 0, 0 self.optimizer.zero_grad() for mb in micro_batches: @@ -848,7 +945,11 @@ def train_step(self, micro_batches: List[Dict[str, Any]]) -> Tuple[float, float] RoutingReplay.clear_all() self._sync_sp_gradients() - + if self._use_distsignsgd and active_microbatches > 0: + scale_model_gradients( + self.model, + get_distsign_grad_scale_factor(active_voter_total), + ) grad_norm = self._clip_and_step() self._maybe_merge_lora() total_loss, grad_norm = self._reduce_metrics(total_loss, grad_norm) @@ -866,6 +967,13 @@ def _count_valid_tokens(self, micro_batches: List[Dict[str, Any]]) -> torch.Tens group=self.ps.fsdp_group if self.pp_enabled else None, ) + def _count_active_microbatches(self, micro_batches: List[Dict[str, Any]]) -> tuple[int, int]: + """Return ``(active_microbatches, active_voter_total)`` over the DP group.""" + return count_active_microbatches( + micro_batches, + group=self.ps.fsdp_group if self.pp_enabled else None, + ) + def _pad_micro_batches_for_pp(self, micro_batches: List[Dict[str, Any]]) -> None: pad_micro_batches_for_pp( micro_batches, @@ -876,7 +984,11 @@ def _pad_micro_batches_for_pp(self, micro_batches: List[Dict[str, Any]]) -> None def _sync_sp_gradients(self) -> None: """All-reduce gradients for CP/Ulysses dims not folded into FSDP.""" - sync_sp_gradients(self.model, self.ps.sp_grad_sync_group) + sync_sp_gradients( + self.model, + self.ps.sp_grad_sync_group, + skip_dtensor_grads=self._use_distsignsgd, + ) def _reduce_metrics(self, total_loss: float, grad_norm: float) -> Tuple[float, float]: """All-reduce loss and grad_norm across DP for logging.""" @@ -921,7 +1033,7 @@ def _forward_backward( outputs.last_hidden_state, loss_fn_name=None, loss_fn_inputs={"labels": labels}, - loss_fn_params=None, + loss_fn_params=self._causallm_loss_params, logits_to_keep=0, ) loss = result.loss @@ -984,12 +1096,8 @@ def _forward_backward_pp( pp_group=self.ps.pp_group, ) gvt = global_valid_tokens.item() - if gvt > 0: - scale = 1.0 / gvt - for model_part in self.model_parts: - for p in model_part.parameters(): - if p.grad is not None: - p.grad.mul_(scale) + if gvt > 0 and not self._use_distsignsgd: + scale_model_gradients(self.model_parts, 1.0 / float(gvt)) return raw_loss / gvt if gvt > 0 else 0.0 def _clip_and_step(self) -> float: @@ -997,9 +1105,13 @@ def _clip_and_step(self) -> float: Returns grad_norm (scalar). """ + clip_value = get_effective_grad_clip_value( + self.args.train.max_grad_norm, + use_distsignsgd=self._use_distsignsgd, + ) grad_norm = clip_gradients( self.model, - self.args.train.max_grad_norm, + clip_value, pp_enabled=self.pp_enabled, pp_group=self.ps.pp_group if self.pp_enabled else None, ) @@ -1139,6 +1251,15 @@ def _finalize(self, total_loss: float, grad_norm: float, lr: float) -> None: synchronize() + # Report peak GPU memory on rank 0 — parseable by benchmark scripts. + if logger.isEnabledFor(10): # DEBUG + device = get_torch_device() + peak_alloc_gb = device.max_memory_allocated() / (1024**3) + peak_reserved_gb = device.max_memory_reserved() / (1024**3) + logger.debug_rank0( + f"[PEAK_MEMORY] peak_alloc_gb={peak_alloc_gb:.3f} peak_reserved_gb={peak_reserved_gb:.3f}" + ) + # Gather full model state for HF save is_lora_training = args.lora.enable_lora or args.lora.enable_qlora save_peft_adapter = is_lora_training and args.lora.merge_lora_interval == 0 diff --git a/src/xorl/trainers/training_utils.py b/src/xorl/trainers/training_utils.py index 2a85cc50..dd6591b4 100644 --- a/src/xorl/trainers/training_utils.py +++ b/src/xorl/trainers/training_utils.py @@ -19,7 +19,18 @@ from xorl.utils.device import get_device_type -def sync_sp_gradients(model: torch.nn.Module, sp_grad_sync_group) -> None: +try: + from torch.distributed._tensor import DTensor +except ImportError: # pragma: no cover - torch 2.10+ always provides DTensor here + DTensor = None + + +def sync_sp_gradients( + model: torch.nn.Module, + sp_grad_sync_group, + *, + skip_dtensor_grads: bool = False, +) -> None: """All-reduce gradients for ring/Ulysses dims not folded into FSDP. SP ranks hold complementary (non-overlapping) parts of the same sequence, @@ -28,11 +39,18 @@ def sync_sp_gradients(model: torch.nn.Module, sp_grad_sync_group) -> None: cp_fsdp_mode="all": group is None → no-op cp_fsdp_mode="ulysses_only": group is ring group cp_fsdp_mode="none": group is unified SP group + + When DistSignSGD is active, FSDP-managed grads perform the exact SP sum + inside the custom reduce-scatter hook before `sign()`. In that case, the + later external SP sync should only touch non-FSDP grads. """ if sp_grad_sync_group is not None: for p in model.parameters(): - if p.grad is not None: - dist.all_reduce(p.grad, op=dist.ReduceOp.SUM, group=sp_grad_sync_group) + if p.grad is None: + continue + if skip_dtensor_grads and DTensor is not None and isinstance(p.grad, DTensor): + continue + dist.all_reduce(p.grad, op=dist.ReduceOp.SUM, group=sp_grad_sync_group) def clip_gradients( @@ -70,6 +88,40 @@ def clip_gradients( return grad_norm +def get_effective_grad_clip_value(max_grad_norm: float, *, use_distsignsgd: bool) -> float: + """Return the clipping threshold to use for the current optimizer path. + + DistSignSGD turns gradients into sign-vote accumulators before the training + loop reaches grad clipping. Clipping those sign votes changes the update + scale by orders of magnitude, so we pass float("inf") to disable clipping + and let the downstream `clip_gradients` call return the unclipped L2 norm + purely for observability. + + Note for log readers: under DistSignSGD the value reported as "grad_norm" + is really the L2 norm of accumulated sign votes (think `vote_l2_norm`), + not a true gradient magnitude — its scale tracks `sqrt(num_params)` and + voter agreement, not the underlying loss landscape. + """ + if use_distsignsgd: + return float("inf") + return max_grad_norm + + +def get_distsign_grad_scale_factor(active_voter_total: int) -> float: + """Return the scale factor that converts accumulated sign votes to a mean. + + `active_voter_total` is the total number of (microbatch, rank) pairs that + actually cast a sign vote — i.e. ranks whose microbatch had at least one + valid token. Ranks with zero valid tokens contribute sign(0) = 0, not a + ±1 vote, so multiplying `active_microbatches * dp_size` would over-count + abstainers and bias the per-step update toward zero on uneven token + distributions. + """ + if active_voter_total <= 0: + return 1.0 + return 1.0 / float(active_voter_total) + + def count_valid_tokens( micro_batches: List[Dict[str, Any]], group=None, @@ -88,6 +140,57 @@ def count_valid_tokens( return global_valid_tokens +def count_active_microbatches( + micro_batches: List[Dict[str, Any]], + group=None, +) -> tuple[int, int]: + """Return ``(active_microbatches, active_voter_total)`` for sign-vote aggregation. + + A single batched all-reduce (op=SUM) is issued for the whole accumulation step: + + - ``active_microbatches``: number of micro-batches in which *any* rank in + ``group`` had at least one valid token. + - ``active_voter_total``: sum over micro-batches of the number of ranks + with valid tokens. This equals the number of (micro-batch, rank) pairs + that contribute a real ±1 sign vote (ranks with zero valid tokens emit + sign(0) = 0 and abstain). + + Callers should use ``active_voter_total`` as the divisor when normalizing + accumulated sign votes; using ``active_microbatches * dp_size`` would + over-count abstainers when token distribution is uneven. + """ + if not micro_batches: + return 0, 0 + + flags = torch.zeros(len(micro_batches), device=get_device_type(), dtype=torch.int64) + for i, mb in enumerate(micro_batches): + labels = mb.get("labels", mb.get("target_tokens")) + if labels is None: + continue + flags[i] = (labels != IGNORE_INDEX).any().to(torch.int64) + dist.all_reduce(flags, op=dist.ReduceOp.SUM, group=group) + active_voter_total = int(flags.sum().item()) + active_microbatches = int((flags > 0).sum().item()) + return active_microbatches, active_voter_total + + +def scale_model_gradients(model_or_models, scale: float) -> None: + """Scale gradients in-place while preserving DTensor metadata.""" + if scale == 1.0: + return + + modules = model_or_models if isinstance(model_or_models, (list, tuple)) else [model_or_models] + seen: set[int] = set() + for module in modules: + for param in module.parameters(): + param_id = id(param) + if param_id in seen: + continue + seen.add(param_id) + if param.grad is not None: + param.grad.mul_(scale) + + def reset_lora_optimizer_states( model: torch.nn.Module, optimizer: torch.optim.Optimizer, @@ -167,13 +270,11 @@ def negotiate_pp_seq_len(micro_batches: List[Dict[str, Any]], pp_group) -> int: return int(t.item()) -@torch.compile -def pp_loss_fn(pred, labels): - """Compiled PP cross-entropy loss (raw CE sum, unnormalized). +def _pp_ce_sum(pred, labels): + """Raw PP cross-entropy sum over all non-ignored tokens (unnormalized). - Returns CE_sum over all non-ignored tokens. Callers are responsible - for dividing gradients by global_valid_tokens after the backward - (either immediately or deferred to optim_step). + Callers are responsible for dividing gradients by global_valid_tokens + after the backward (either immediately or deferred to optim_step). """ return F.cross_entropy( pred.flatten(0, 1).float(), @@ -183,6 +284,23 @@ def pp_loss_fn(pred, labels): ) +_pp_ce_sum_compiled = torch.compile(_pp_ce_sum) + + +def make_pp_loss_fn(ce_mode: str = "compiled"): + """Return the PP cross-entropy loss variant selected by ``ce_mode``. + + 'compiled' (default) returns the torch.compile'd CE sum; 'eager' + returns the uncompiled baseline (useful for debugging or when compile + regresses). + """ + if ce_mode == "eager": + return _pp_ce_sum + if ce_mode == "compiled": + return _pp_ce_sum_compiled + raise ValueError(f"Unknown ce_mode: {ce_mode!r} (expected 'eager' or 'compiled')") + + def pad_micro_batches_for_pp( micro_batches: List[Dict[str, Any]], sample_packing_sequence_len: int, diff --git a/src/xorl/utils/count_flops.py b/src/xorl/utils/count_flops.py index 214779c0..ceccf330 100644 --- a/src/xorl/utils/count_flops.py +++ b/src/xorl/utils/count_flops.py @@ -114,6 +114,7 @@ def __init__( "qwen3": self._estimate_qwen2_flops, "xorl_qwen3_5": self._estimate_qwen3_5_flops, "xorl_qwen3_5_moe": self._estimate_qwen3_5_moe_flops, + "glm4_moe": self._estimate_glm4_moe_flops, } self.config = config @@ -347,6 +348,68 @@ def _estimate_qwen3_5_moe_flops(self, tokens_sum, batch_seqlens, delta_time): return (dense_N_flops + attn_qkv_flops) / delta_time / 1e12 + def _estimate_glm4_moe_flops(self, tokens_sum, batch_seqlens, delta_time): + """FLOPs estimate for GLM-4 MoE (e.g. GLM-4.7). + + Architecture: GQA attention + MoE FFN with shared experts. + First ``first_k_dense_replace`` layers use a dense MLP; the remaining + layers use routed experts plus a shared expert. + """ + hidden_size = self.config.hidden_size + vocab_size = self.config.vocab_size + num_hidden_layers = self.config.num_hidden_layers + num_attention_heads = self.config.num_attention_heads + num_key_value_heads = self.config.num_key_value_heads + moe_intermediate_size = self.config.moe_intermediate_size + moe_num_expert = self.config.n_routed_experts + moe_topk = self.config.num_experts_per_tok + n_shared_experts = getattr(self.config, "n_shared_experts", 0) + first_k_dense_replace = getattr(self.config, "first_k_dense_replace", 0) + dense_intermediate_size = self.config.intermediate_size # MLP dim for dense layers + + m = self._m + head_dim = getattr(self.config, "head_dim", hidden_size // num_attention_heads) + q_size = num_attention_heads * head_dim + k_size = num_key_value_heads * head_dim + v_size = num_key_value_heads * head_dim + + # Attention linear projections (Q, K, V, O) — same for all layers + attn_linear_N = hidden_size * (q_size + k_size + v_size + num_attention_heads * head_dim) + + # MoE layer FLOPs: router + routed experts (top-k) + shared expert(s) + 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 + # Shared expert uses moe_intermediate_size (fused gate_up + down) + shared_gate_up_N = hidden_size * moe_intermediate_size * n_shared_experts * 2 + shared_down_N = hidden_size * moe_intermediate_size * n_shared_experts + + moe_layer_flops = ( + m["attn_linear"] * attn_linear_N * tokens_sum + + m["router"] * router_N * tokens_sum + + m["gate"] * gate_up_N * tokens_sum + + m["down"] * down_N * tokens_sum + + m["gate"] * shared_gate_up_N * tokens_sum + + m["down"] * shared_down_N * tokens_sum + ) + + # Dense layer FLOPs (first K layers): standard SwiGLU MLP (gate_up_proj + down_proj = 3× H×I) + dense_mlp_N = hidden_size * dense_intermediate_size * 3 + dense_layer_flops = m["attn_linear"] * attn_linear_N * tokens_sum + m["dense_mlp"] * dense_mlp_N * tokens_sum + + num_moe_layers = num_hidden_layers - first_k_dense_replace + embed_lm_N = vocab_size * hidden_size # lm_head only; embedding is a lookup + + dense_N_flops = ( + moe_layer_flops * num_moe_layers + dense_layer_flops * first_k_dense_replace + 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 + + return (dense_N_flops + attn_qkv_flops) / delta_time / 1e12 + def _estimate_qwen2_flops(self, tokens_sum, batch_seqlens, delta_time): hidden_size = self.config.hidden_size vocab_size = self.config.vocab_size diff --git a/tests/_helpers/__init__.py b/tests/_helpers/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/_helpers/opd.py b/tests/_helpers/opd.py new file mode 100644 index 00000000..264e9400 --- /dev/null +++ b/tests/_helpers/opd.py @@ -0,0 +1,126 @@ +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +from typing import Mapping + +import torch +import torch.nn.functional as F +from safetensors.torch import save_file + + +@dataclass(frozen=True) +class TeacherFiles: + heads: dict[str, str] + hidden_caches: dict[str, str] + + +def save_tensor_file(path: str | Path, key: str, tensor: torch.Tensor) -> str: + path = Path(path) + save_file({key: tensor.detach().cpu().contiguous()}, str(path)) + return str(path) + + +def make_teacher_files( + tmp_path: Path, + teacher_heads: Mapping[str, torch.Tensor], + teacher_hidden_caches: Mapping[str, torch.Tensor], +) -> TeacherFiles: + head_paths: dict[str, str] = {} + cache_paths: dict[str, str] = {} + for teacher_id in teacher_heads: + head_paths[teacher_id] = save_tensor_file( + tmp_path / f"teacher_{teacher_id}_head.safetensors", + "lm_head.weight", + teacher_heads[teacher_id], + ) + cache_paths[teacher_id] = save_tensor_file( + tmp_path / f"teacher_{teacher_id}_hidden.safetensors", + "hidden_states", + teacher_hidden_caches[teacher_id], + ) + return TeacherFiles(heads=head_paths, hidden_caches=cache_paths) + + +def save_teacher_hidden_cache(hiddens: list[torch.Tensor], path: str | Path) -> list[list[int]]: + """Concatenate per-sample hidden states and return per-sample cache indices.""" + cache_indices: list[list[int]] = [] + offset = 0 + for hidden in hiddens: + cache_indices.append(list(range(offset, offset + hidden.shape[0]))) + offset += hidden.shape[0] + save_tensor_file(path, "hidden_states", torch.cat(hiddens, dim=0)) + return cache_indices + + +def reference_opd_loss( + hidden_states: torch.Tensor, + weight: torch.Tensor, + labels: torch.Tensor, + teacher_hidden_states: torch.Tensor, + teacher_weight: torch.Tensor, + teacher_weights: torch.Tensor | None = None, + *, + ignore_index: int = -100, + normalization_denominator: torch.Tensor | int | float | None = None, +) -> torch.Tensor: + h = hidden_states.reshape(-1, hidden_states.size(-1)) + th = teacher_hidden_states.reshape(-1, teacher_hidden_states.size(-1)) + labels_flat = labels.reshape(-1) + valid = labels_flat != ignore_index + + if not valid.any(): + return h.float().sum() * 0.0 + weight.float().sum() * 0.0 + + student_logits = h[valid] @ weight.t() + teacher_logits = th[valid] @ teacher_weight.t() + student_log_probs = F.log_softmax(student_logits.float(), dim=-1) + teacher_log_probs = F.log_softmax(teacher_logits.float(), dim=-1) + token_kl = (student_log_probs.exp() * (student_log_probs - teacher_log_probs)).sum(dim=-1) + if teacher_weights is not None: + token_kl = token_kl * teacher_weights.reshape(-1)[valid].to(token_kl.device) + + if normalization_denominator is None: + denom = valid.sum().to(device=token_kl.device, dtype=torch.float32).clamp(min=1.0) + elif torch.is_tensor(normalization_denominator): + denom = normalization_denominator.to(device=token_kl.device, dtype=torch.float32).clamp(min=1.0) + else: + denom = torch.tensor(float(normalization_denominator), device=token_kl.device, dtype=torch.float32).clamp( + min=1.0 + ) + return token_kl.sum() / denom + + +def reference_grouped_opd_loss( + batch: Mapping[str, object], + student_hidden_table: torch.Tensor, + student_head: torch.Tensor, + teacher_hidden_caches: Mapping[str, torch.Tensor], + teacher_heads: Mapping[str, torch.Tensor], + *, + ignore_index: int = -100, +) -> torch.Tensor: + input_ids = torch.tensor(batch["input_ids"], dtype=torch.long) + labels = torch.tensor(batch["labels"], dtype=torch.long) + teacher_ids = torch.tensor(batch["teacher_ids"], dtype=torch.long) + cache_indices = torch.tensor(batch["teacher_cache_indices"], dtype=torch.long) + teacher_weights = torch.tensor(batch["teacher_weights"], dtype=torch.float32) + + hidden_states = student_hidden_table[input_ids].reshape(-1, student_hidden_table.shape[-1]) + labels_flat = labels.reshape(-1) + teacher_ids_flat = teacher_ids.reshape(-1) + cache_indices_flat = cache_indices.reshape(-1) + weights_flat = teacher_weights.reshape(-1) + + valid = labels_flat != ignore_index + token_losses = [] + for idx in valid.nonzero(as_tuple=True)[0].tolist(): + teacher_id = str(int(teacher_ids_flat[idx].item())) + student_logits = hidden_states[idx] @ student_head.t() + teacher_hidden = teacher_hidden_caches[teacher_id][cache_indices_flat[idx]] + teacher_logits = teacher_hidden @ teacher_heads[teacher_id].t() + student_log_probs = F.log_softmax(student_logits.float(), dim=-1) + teacher_log_probs = F.log_softmax(teacher_logits.float(), dim=-1) + token_kl = (student_log_probs.exp() * (student_log_probs - teacher_log_probs)).sum() + token_losses.append(token_kl * weights_flat[idx]) + return torch.stack(token_losses).sum() / valid.sum().clamp(min=1) diff --git a/tests/data/collators/test_packing_concat_collator.py b/tests/data/collators/test_packing_concat_collator.py index ae26657f..5b03c0cf 100644 --- a/tests/data/collators/test_packing_concat_collator.py +++ b/tests/data/collators/test_packing_concat_collator.py @@ -127,3 +127,34 @@ def test_extra_fields_and_multiple_seqs(self, mock_parallel_state): batch2 = collator(features2) assert batch2["input_ids"].shape == (1, 9) assert torch.equal(batch2["position_ids"], torch.tensor([[0, 1, 2, 0, 1, 0, 1, 2, 3]])) + + @patch("xorl.data.collators.packing_concat_collator.get_parallel_state") + def test_teacher_hidden_states_concatenate_and_pad_as_sequence_field(self, mock_parallel_state): + mock_ps = Mock() + mock_ps.cp_enabled = False + mock_parallel_state.return_value = mock_ps + collator = PackingConcatCollator(pad_to_multiple_of=4) + features = [ + { + "input_ids": torch.tensor([1, 2], dtype=torch.long), + "attention_mask": torch.tensor([1, 1], dtype=torch.long), + "labels": torch.tensor([2, -100], dtype=torch.long), + "position_ids": torch.tensor([0, 1], dtype=torch.long), + "teacher_hidden_states": torch.tensor([[0.25, 0.5], [1.25, 1.5]]), + }, + { + "input_ids": torch.tensor([3], dtype=torch.long), + "attention_mask": torch.tensor([1], dtype=torch.long), + "labels": torch.tensor([-100], dtype=torch.long), + "position_ids": torch.tensor([0], dtype=torch.long), + "teacher_hidden_states": torch.tensor([[2.25, 2.5]]), + }, + ] + + batch = collator(features) + + assert batch["teacher_hidden_states"].shape == (1, 4, 2) + torch.testing.assert_close( + batch["teacher_hidden_states"], + torch.tensor([[[0.25, 0.5], [1.25, 1.5], [2.25, 2.5], [0.0, 0.0]]]), + ) diff --git a/tests/data/collators/test_sequence_shard_collator.py b/tests/data/collators/test_sequence_shard_collator.py index 216ae5c9..e82d2db6 100644 --- a/tests/data/collators/test_sequence_shard_collator.py +++ b/tests/data/collators/test_sequence_shard_collator.py @@ -156,3 +156,35 @@ def test_sp_splitting_padding_and_flash_attn_kwargs(self, mock_parallel_state): r_single = collator_fa(batch_single) assert r_single["input_ids"].shape[-1] == 5 assert r_single["labels"].shape[-1] == 5 + + @patch("xorl.data.collators.sequence_shard_collator.get_parallel_state") + def test_teacher_hidden_states_are_sharded_with_token_fields(self, mock_parallel_state): + mock_parallel_state.return_value = _make_mock_ps(cp_size=2, cp_rank=1) + collator = TextSequenceShardCollator(pad_token_id=0) + teacher_hidden_states = torch.tensor( + [ + [ + [0.25, 0.5], + [1.25, 1.5], + [2.25, 2.5], + [3.25, 3.5], + [4.25, 4.5], + ] + ] + ) + batch = { + "input_ids": torch.tensor([[1, 2, 3, 4, 5]]), + "attention_mask": torch.tensor([[1, 1, 1, 1, 1]]), + "labels": torch.tensor([[2, 3, 4, 5, IGNORE_INDEX]]), + "position_ids": torch.tensor([[0, 1, 2, 3, 4]]), + "teacher_hidden_states": teacher_hidden_states, + } + + result = collator(batch) + + assert result["input_ids"].shape[-1] == 3 + assert result["teacher_hidden_states"].shape == (1, 3, 2) + torch.testing.assert_close( + result["teacher_hidden_states"], + torch.tensor([[[3.25, 3.5], [4.25, 4.5], [0.0, 0.0]]]), + ) diff --git a/tests/data/prepare/test_shared.py b/tests/data/prepare/test_shared.py index 1421ecd2..dc854f08 100644 --- a/tests/data/prepare/test_shared.py +++ b/tests/data/prepare/test_shared.py @@ -4,12 +4,9 @@ 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 ( - _check_if_hub_dataset, - _get_remote_filesystem, create_train_validation_split, datasets_with_name_generator, get_dataset_type, @@ -82,39 +79,6 @@ def test_expansion_and_passthrough_behaviors(self): assert get_dataset_type(_make_config(path=path)) == expected_type -class TestHubAndRemoteDetection: - """Tests for _check_if_hub_dataset and _get_remote_filesystem.""" - - @patch("xorl.data.prepare.shared.snapshot_download") - def test_hub_detection_and_remote_filesystem(self, mock_snapshot_download): - """Covers hub dataset detection (valid/invalid) and remote filesystem resolution.""" - # Valid hub dataset - mock_snapshot_download.return_value = "/path/to/dataset" - assert _check_if_hub_dataset(_make_config(path="user/ds"), use_auth_token=False) is True - - # Invalid hub dataset - - mock_response = Mock() - mock_response.status_code = 404 - mock_response.headers = {} - mock_snapshot_download.side_effect = RepositoryNotFoundError("not found", response=mock_response) - assert _check_if_hub_dataset(_make_config(path="invalid/ds"), use_auth_token=False) is False - - # Non-remote path - fs, opts = _get_remote_filesystem("local/path") - assert fs is None - assert opts == {} - - # S3 path - pytest.importorskip("s3fs") - with patch("s3fs.S3FileSystem") as mock_s3fs_class: - mock_fs = Mock() - mock_s3fs_class.return_value = mock_fs - fs, opts = _get_remote_filesystem("s3://bucket/path") - assert fs == mock_fs - assert opts == {"anon": False} - - class TestSplitAndMerge: """Tests for create_train_validation_split and merge_datasets.""" diff --git a/tests/distributed/test_bf16_a2a_fsdp_hook.py b/tests/distributed/test_bf16_a2a_fsdp_hook.py new file mode 100644 index 00000000..187d70f4 --- /dev/null +++ b/tests/distributed/test_bf16_a2a_fsdp_hook.py @@ -0,0 +1,128 @@ +"""End-to-end smoke test: BF16StochasticAllToAllReduceScatter wired into FSDP2. + +Verifies that ``FSDPModule.set_custom_reduce_scatter`` accepts our custom +``ReduceScatter`` and that backward + optimizer step produces non-NaN updated +weights. Compares against FP32 reduce-scatter on the same model and grads; +loss should be close (within a generous tolerance given BF16 transit). +""" + +from __future__ import annotations + +import os +import sys +from pathlib import Path + +import pytest +import torch +import torch.distributed as dist +import torch.nn as nn +from torch.distributed._composable.fsdp import MixedPrecisionPolicy, fully_shard + +from xorl.distributed.fsdp2 import BF16StochasticAllToAllReduceScatter +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] + + +def _setup_dist() -> torch.device: + local_rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(local_rank) + dist.init_process_group(backend=get_nccl_backend()) + return torch.device("cuda", local_rank) + + +class TinyMLP(nn.Module): + def __init__(self, hidden=64, intermediate=128): + super().__init__() + self.fc1 = nn.Linear(hidden, intermediate, bias=False) + self.fc2 = nn.Linear(intermediate, hidden, bias=False) + + def forward(self, x): + return self.fc2(torch.nn.functional.gelu(self.fc1(x))) + + +def _step(model, x, y): + out = model(x) + loss = (out - y).pow(2).mean() + loss.backward() + return float(loss.detach().item()) + + +def _run() -> None: + device = _setup_dist() + torch.manual_seed(123 + dist.get_rank()) + + hidden, intermediate = 64, 128 + x = torch.randn(8, hidden, device=device, dtype=torch.bfloat16) + y = torch.randn(8, hidden, device=device, dtype=torch.bfloat16) + + mp_policy = MixedPrecisionPolicy(param_dtype=torch.bfloat16, reduce_dtype=torch.float32) + mesh = dist.device_mesh.init_device_mesh("cuda", (dist.get_world_size(),), mesh_dim_names=("dp_shard",)) + + # Two identical models seeded the same; one with FP32 reduce-scatter, one with BF16-a2a. + torch.manual_seed(7) + model_ref = TinyMLP(hidden, intermediate).to(device) + torch.manual_seed(7) + model_test = TinyMLP(hidden, intermediate).to(device) + + fully_shard(model_ref.fc1, mesh=mesh, mp_policy=mp_policy) + fully_shard(model_ref.fc2, mesh=mesh, mp_policy=mp_policy) + fully_shard(model_ref, mesh=mesh, mp_policy=mp_policy) + + fully_shard(model_test.fc1, mesh=mesh, mp_policy=mp_policy) + fully_shard(model_test.fc2, mesh=mesh, mp_policy=mp_policy) + fully_shard(model_test, mesh=mesh, mp_policy=mp_policy) + # Install the custom reduce-scatter on both Linear FSDP units + model_test.fc1.set_custom_reduce_scatter(BF16StochasticAllToAllReduceScatter()) + model_test.fc2.set_custom_reduce_scatter(BF16StochasticAllToAllReduceScatter()) + + loss_ref = _step(model_ref, x, y) + loss_test = _step(model_test, x, y) + + # Compare grads in fc1 / fc2 (DTensor → local_tensor) + g1_ref = model_ref.fc1.weight.grad._local_tensor + g1_test = model_test.fc1.weight.grad._local_tensor + g2_ref = model_ref.fc2.weight.grad._local_tensor + g2_test = model_test.fc2.weight.grad._local_tensor + + rel1 = (g1_test - g1_ref).abs().max() / (g1_ref.abs().max() + 1e-6) + rel2 = (g2_test - g2_ref).abs().max() / (g2_ref.abs().max() + 1e-6) + + if dist.get_rank() == 0: + print(f"loss ref={loss_ref:.6f}, test={loss_test:.6f}") + print(f"fc1 max-rel-err: {rel1.item():.4e}") + print(f"fc2 max-rel-err: {rel2.item():.4e}") + + # Both gradients must be finite and within a generous BF16 tolerance. + assert torch.isfinite(g1_test).all() and torch.isfinite(g2_test).all(), ( + "BF16 a2a path produced non-finite gradients" + ) + assert rel1.item() < 0.05, f"fc1 grad relative error {rel1.item()} too large" + assert rel2.item() < 0.05, f"fc2 grad relative error {rel2.item()} too large" + + dist.barrier() + dist.destroy_process_group() + + +def _main() -> None: + _run() + + +if __name__ != "__main__": + + @skip_if_gpu_count_less_than(2) + def test_bf16_a2a_reduce_scatter_runs_inside_fsdp2(): + result = run_distributed_script(__file__, num_gpus=2, timeout=180) + result.assert_success("BF16 a2a reduce-scatter should integrate cleanly with fully_shard") + + +if __name__ == "__main__": + _main() diff --git a/tests/distributed/test_bf16_a2a_reduce.py b/tests/distributed/test_bf16_a2a_reduce.py new file mode 100644 index 00000000..7318c557 --- /dev/null +++ b/tests/distributed/test_bf16_a2a_reduce.py @@ -0,0 +1,141 @@ +"""Distributed correctness tests for ``BF16StochasticAllToAllReduceScatter``. + +Verifies the custom reduce-scatter (stochastic-round FP32→BF16, all-to-all, +local FP32 sum) produces results numerically close to native FP32 +reduce-scatter, with bias-in-expectation near zero. +""" + +from __future__ import annotations + +import os +import sys +from pathlib import Path + +import pytest +import torch +import torch.distributed as dist +from torch.distributed.distributed_c10d import ReduceOp + +from xorl.distributed.fsdp2 import BF16StochasticAllToAllReduceScatter +from xorl.distributed.fsdp2.bf16_a2a_reduce import _canonical_reduce_op +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] + + +@pytest.mark.cpu +def test_canonical_reduce_op_accepts_fsdp_wrapped_ops(): + assert _canonical_reduce_op(ReduceOp(ReduceOp.SUM)) == dist.ReduceOp.SUM + assert _canonical_reduce_op(ReduceOp(ReduceOp.AVG)) == dist.ReduceOp.AVG + assert _canonical_reduce_op(dist.ReduceOp.SUM) == dist.ReduceOp.SUM + assert _canonical_reduce_op(dist._make_nccl_premul_sum(1.0)) == dist.ReduceOp.SUM + + +def _world_size() -> int: + return int(os.environ["WORLD_SIZE"]) + + +def _local_rank() -> int: + return int(os.environ["LOCAL_RANK"]) + + +def _setup_dist() -> torch.device: + local_rank = _local_rank() + torch.cuda.set_device(local_rank) + dist.init_process_group(backend=get_nccl_backend()) + return torch.device("cuda", local_rank) + + +def _run() -> None: + device = _setup_dist() + rank = dist.get_rank() + world = dist.get_world_size() + + # Each rank generates a chunk of FP32 grad. We want to reduce-scatter: + # the input on every rank is the FULL flat unsharded grad, viewed as + # ``world * chunk`` elements; reduce-scatter sums across ranks and + # gives each rank its ``chunk_numel`` slice of the global sum. + chunk_numel = 4096 + total_numel = chunk_numel * world + + torch.manual_seed(0xCAFE + rank) + # Per-rank gradient (independent across ranks). + local_grad = torch.randn(total_numel, dtype=torch.float32, device=device) + + # ---- Reference: native FP32 reduce-scatter ---- + ref_out = torch.empty(chunk_numel, dtype=torch.float32, device=device) + dist.reduce_scatter_tensor(ref_out, local_grad.clone(), op=dist.ReduceOp.SUM) + + # ---- Test: BF16 stochastic-rounded a2a + FP32 local sum ---- + comm = BF16StochasticAllToAllReduceScatter() + test_out = comm.allocate((chunk_numel,), dtype=torch.float32, device=device) + comm(test_out, local_grad.clone(), group=dist.group.WORLD, op=dist.ReduceOp.SUM) + + # ---- Bound the per-element error ---- + # Each rank's contribution is stochastically rounded FP32→BF16 with at most + # one ulp of noise. After summing ``world`` such contributions, the error + # is bounded by sum of |x_r| * 2^-7 in the worst case. Compute this bound. + abs_input = local_grad.abs() + err_bound_local = abs_input * (2**-7) # per-rank max error envelope + # Get this rank's slice of the global error bound — match what reduce-scatter does. + err_bound_full_sum = err_bound_local.clone() + dist.all_reduce(err_bound_full_sum, op=dist.ReduceOp.SUM) + # Slice this rank's chunk of the bound. + bound_chunks = err_bound_full_sum.view(world, chunk_numel) + bound_for_my_chunk = bound_chunks[rank] + + abs_err = (test_out - ref_out).abs() + # Max element-wise should be <= our bound, with some headroom for FP32 + # rounding in the local sum. Use 4x headroom. + max_err = abs_err.max().item() + max_bound = bound_for_my_chunk.max().item() * 4 + 1e-6 + assert max_err < max_bound, f"[rank {rank}] BF16 a2a max err {max_err:.4e} exceeds bound {max_bound:.4e}" + + # Bias-in-expectation: average over many trials should approach the FP32 reference. + # Use the same input tensor; only the stochastic rounding noise differs. + n_trials = 200 + accum = torch.zeros_like(test_out) + for _ in range(n_trials): + out = comm.allocate((chunk_numel,), dtype=torch.float32, device=device) + comm(out, local_grad.clone(), group=dist.group.WORLD, op=dist.ReduceOp.SUM) + accum += out + mean = accum / n_trials + mean_err = (mean - ref_out).abs().max().item() + # Standard error of the mean ~ bound / sqrt(n_trials). For n=200 and BF16 + # bound ~|x|/128, SEM ~ |x| * 1e-3. Allow generous 5x headroom. + sem_bound = max_bound / (n_trials**0.5) * 5 + assert mean_err < sem_bound, ( + f"[rank {rank}] BF16 a2a is biased: mean err over {n_trials} trials = " + f"{mean_err:.4e}, expected < {sem_bound:.4e}" + ) + + if rank == 0: + print(f"[rank 0] BF16 a2a max err = {max_err:.4e}, bound = {max_bound:.4e}") + print(f"[rank 0] BF16 a2a unbiased mean err over {n_trials} trials = {mean_err:.4e}") + + dist.barrier() + dist.destroy_process_group() + + +def _main() -> None: + _run() + + +if __name__ != "__main__": + + @skip_if_gpu_count_less_than(4) + def test_bf16_a2a_reduce_scatter_matches_fp32_within_bound(): + result = run_distributed_script(__file__, num_gpus=4, timeout=180) + result.assert_success("BF16 a2a reduce-scatter should match FP32 within BF16 ulp bound") + + +if __name__ == "__main__": + _main() diff --git a/tests/distributed/test_deepep_async_combine_guard.py b/tests/distributed/test_deepep_async_combine_guard.py new file mode 100644 index 00000000..fd809837 --- /dev/null +++ b/tests/distributed/test_deepep_async_combine_guard.py @@ -0,0 +1,55 @@ +from types import SimpleNamespace + +import pytest +import torch + +from xorl.distributed.moe import deepep + + +pytestmark = pytest.mark.cpu + + +def test_deepep_async_combine_is_synchronous_by_default(monkeypatch): + captured = {} + + def fake_apply(expert_output, buffer, ctx, async_combine): + del buffer, ctx + captured["async_combine"] = async_combine + return expert_output + + monkeypatch.setattr(deepep, "_ALLOW_UNSAFE_ASYNC_COMBINE", False) + monkeypatch.setattr(deepep._FusedUnpermuteAndCombine, "apply", staticmethod(fake_apply)) + + expert_output = torch.ones(1, 2) + result = deepep.tokens_post_combine( + buffer=None, + expert_output=expert_output, + ctx=SimpleNamespace(), + async_combine=True, + ) + + assert result is expert_output + assert captured["async_combine"] is False + + +def test_deepep_async_combine_can_be_unsafely_opted_in(monkeypatch): + captured = {} + + def fake_apply(expert_output, buffer, ctx, async_combine): + del buffer, ctx + captured["async_combine"] = async_combine + return expert_output + + monkeypatch.setattr(deepep, "_ALLOW_UNSAFE_ASYNC_COMBINE", True) + monkeypatch.setattr(deepep._FusedUnpermuteAndCombine, "apply", staticmethod(fake_apply)) + + expert_output = torch.ones(1, 2) + result = deepep.tokens_post_combine( + buffer=None, + expert_output=expert_output, + ctx=SimpleNamespace(), + async_combine=True, + ) + + assert result is expert_output + assert captured["async_combine"] is True diff --git a/tests/distributed/test_deepseek_v3_ep_checkpoint.py b/tests/distributed/test_deepseek_v3_ep_checkpoint.py new file mode 100644 index 00000000..8d1b885c --- /dev/null +++ b/tests/distributed/test_deepseek_v3_ep_checkpoint.py @@ -0,0 +1,163 @@ +"""Distributed EP checkpoint slicing smoke test for DeepSeek/Kimi.""" + +import math +import os + +import torch +import torch.distributed as dist + +from xorl.models.transformers.deepseek_v3.checkpoint_handler import DeepseekV3CheckpointHandler + + +NUM_EXPERTS = 4 + + +def _expert_weight(expert_idx: int, proj: str) -> torch.Tensor: + hidden_size = 2 + intermediate_size = 3 + value = float(expert_idx * 10 + {"gate": 1, "up": 2, "down": 3}[proj]) + if proj == "down": + return torch.full((hidden_size, intermediate_size), value) + return torch.full((intermediate_size, hidden_size), value) + + +def _pack_int4(values: torch.Tensor) -> torch.Tensor: + if values.dtype != torch.int8: + raise ValueError(f"Expected int8 values to pack, got {values.dtype}") + if values.ndim != 2: + raise ValueError(f"Expected rank-2 tensor to pack, got {tuple(values.shape)}") + + num_bits = 4 + pack_factor = 32 // num_bits + unsigned = (values + (1 << (num_bits - 1))).to(torch.uint8) + pad_cols = (-values.shape[1]) % pack_factor + if pad_cols: + unsigned = torch.nn.functional.pad(unsigned, (0, pad_cols)) + reshaped = unsigned.view(values.shape[0], -1, pack_factor).to(torch.int32) + bit_shifts = torch.arange(pack_factor, dtype=torch.int32) * num_bits + return (reshaped << bit_shifts).sum(dim=2, dtype=torch.int32) + + +def _packed_expert_weight(expert_idx: int, proj: str) -> dict[str, torch.Tensor]: + dense_weight = _expert_weight(expert_idx, proj) + quantized = torch.ones_like(dense_weight, dtype=torch.int8) + num_groups = max(1, math.ceil(dense_weight.shape[1] / 32)) + scales = torch.full((dense_weight.shape[0], num_groups), dense_weight.flatten()[0].item(), dtype=torch.float32) + return { + "weight_packed": _pack_int4(quantized), + "weight_scale": scales, + "weight_shape": torch.tensor(dense_weight.shape, dtype=torch.int64), + } + + +def _setup(): + dist.init_process_group(backend="gloo") + return dist.get_rank(), dist.get_world_size() + + +def _load_local_expert_slice(rank: int, world_size: int) -> dict: + handler = DeepseekV3CheckpointHandler(num_experts=NUM_EXPERTS, ep_rank=rank, ep_size=world_size) + skip_key = handler.get_skip_key_fn() + loaded = {} + + for expert_idx in range(NUM_EXPERTS): + for proj in ("gate", "up", "down"): + key = f"language_model.model.layers.0.mlp.experts.{expert_idx}.{proj}_proj.weight" + if skip_key is not None and skip_key(key): + loaded.update(handler.on_skip_weight(key)) + else: + loaded.update(handler.on_load_weight(key, _expert_weight(expert_idx, proj))) + + gate_up = loaded["model.layers.0.mlp.experts.gate_up_proj"] + down = loaded["model.layers.0.mlp.experts.down_proj"] + + return { + "rank": rank, + "gate_up_shape": tuple(gate_up.shape), + "down_shape": tuple(down.shape), + "gate_ids": gate_up[:, 0, 0].tolist(), + "down_ids": down[:, 0, 0].tolist(), + } + + +def _load_local_packed_expert_slice(rank: int, world_size: int) -> dict: + handler = DeepseekV3CheckpointHandler(num_experts=NUM_EXPERTS, ep_rank=rank, ep_size=world_size) + skip_key = handler.get_skip_key_fn() + loaded = {} + + for expert_idx in range(NUM_EXPERTS): + for proj in ("gate", "up", "down"): + for suffix, tensor in _packed_expert_weight(expert_idx, proj).items(): + key = f"language_model.model.layers.0.mlp.experts.{expert_idx}.{proj}_proj.{suffix}" + if skip_key is not None and skip_key(key): + loaded.update(handler.on_skip_weight(key)) + else: + loaded.update(handler.on_load_weight(key, tensor)) + + gate_up = loaded["model.layers.0.mlp.experts.gate_up_proj"] + down = loaded["model.layers.0.mlp.experts.down_proj"] + + return { + "rank": rank, + "gate_up_shape": tuple(gate_up.shape), + "down_shape": tuple(down.shape), + "gate_ids": gate_up[:, 0, 0].tolist(), + "down_ids": down[:, 0, 0].tolist(), + } + + +def main(): + rank, world_size = _setup() + assert world_size == 2, f"Expected 2 ranks, got {world_size}" + + summary = _load_local_expert_slice(rank, world_size) + packed_summary = _load_local_packed_expert_slice(rank, world_size) + gathered = [None] * world_size + packed_gathered = [None] * world_size + dist.all_gather_object(gathered, summary) + dist.all_gather_object(packed_gathered, packed_summary) + + if rank == 0: + gathered = sorted(gathered, key=lambda item: item["rank"]) + assert gathered[0]["gate_up_shape"] == (2, 2, 6) + assert gathered[0]["down_shape"] == (2, 3, 2) + assert gathered[0]["gate_ids"] == [1.0, 11.0] + assert gathered[0]["down_ids"] == [3.0, 13.0] + assert gathered[1]["gate_ids"] == [21.0, 31.0] + assert gathered[1]["down_ids"] == [23.0, 33.0] + + packed_gathered = sorted(packed_gathered, key=lambda item: item["rank"]) + assert packed_gathered[0]["gate_up_shape"] == (2, 2, 6) + assert packed_gathered[0]["down_shape"] == (2, 3, 2) + assert packed_gathered[0]["gate_ids"] == [1.0, 11.0] + assert packed_gathered[0]["down_ids"] == [3.0, 13.0] + assert packed_gathered[1]["gate_ids"] == [21.0, 31.0] + assert packed_gathered[1]["down_ids"] == [23.0, 33.0] + print("DeepseekV3 EP checkpoint slicing passed") + + dist.barrier() + dist.destroy_process_group() + + +if __name__ != "__main__": + import pytest + + from tests.distributed.distributed_utils import run_distributed_script + + SCRIPT_PATH = os.path.abspath(__file__) + REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")) + + @pytest.mark.cpu + @pytest.mark.distributed + def test_deepseek_v3_ep_checkpoint_2proc(): + result = run_distributed_script( + SCRIPT_PATH, + num_gpus=2, + timeout=120, + extra_env={"PYTHONPATH": os.path.join(REPO_ROOT, "src")}, + ) + result.assert_success() + + +if __name__ == "__main__": + main() diff --git a/tests/distributed/test_ep_lora_weight_slicing.py b/tests/distributed/test_ep_lora_weight_slicing.py index 8dcef965..67c3d4d0 100644 --- a/tests/distributed/test_ep_lora_weight_slicing.py +++ b/tests/distributed/test_ep_lora_weight_slicing.py @@ -99,7 +99,7 @@ def test_ep_plan_and_shard_tensor(self): lora_tensor = torch.randn(global_shape) for ep_rank, expected_slice in [(0, slice(0, 4)), (1, slice(4, 8))]: - with patch("xorl.distributed.parallel_state.get_parallel_state") as mock_ps: + with patch("xorl.distributed.parallel_plan.get_parallel_state") as mock_ps: mock_state = MagicMock() mock_state.ep_enabled = True mock_state.ep_rank = ep_rank diff --git a/tests/distributed/test_loss_metric_reductions.py b/tests/distributed/test_loss_metric_reductions.py new file mode 100644 index 00000000..737bc8c1 --- /dev/null +++ b/tests/distributed/test_loss_metric_reductions.py @@ -0,0 +1,225 @@ +"""Real-NCCL tests for the IS-metric cross-rank reduction primitives. + +Targets the three helpers in ``xorl.server.runner.model_runner`` that +``dist.all_reduce`` IS metrics across process groups: + +- ``_sp_allreduce_kl_metrics`` — per-mb CP/Ulysses reduction. +- ``ModelRunner._accumulate_is_metrics`` — cross-mb accumulation. +- ``ModelRunner._finalize_is_metrics`` — cross-DP reduction + finalization. +""" + +from __future__ import annotations + +import math +import os +from types import SimpleNamespace +from unittest.mock import patch + +import pytest +import torch +import torch.distributed as dist +from distributed_utils import run_distributed_script, skip_if_gpu_count_less_than + +from xorl.server.runner import model_runner as mr +from xorl.utils.device import get_nccl_backend + + +pytestmark = [pytest.mark.distributed] + + +def _setup_dist() -> torch.device: + local_rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(local_rank) + dist.init_process_group(backend=get_nccl_backend()) + return torch.device("cuda", local_rank) + + +def _make_metrics(device: torch.device, **values) -> dict: + """Plain scalars → the {valid_tokens: int, others: float64-tensor} shape + that the IS-metric helpers consume in production.""" + return { + k: v if k == "valid_tokens" else torch.as_tensor(v, dtype=torch.float64, device=device) + for k, v in values.items() + } + + +def _finalize_with_world_dp(accumulated: dict, result: dict) -> None: + """``_finalize_is_metrics`` resolves the DP group via ``get_parallel_state()``. + Spinning up the full mesh for a unit test is overkill, so stub it to + treat WORLD as the DP group for the duration of the call.""" + ps = SimpleNamespace(dp_enabled=True, dp_group=dist.group.WORLD) + with patch.object(mr, "get_parallel_state", lambda: ps): + mr.ModelRunner._finalize_is_metrics(accumulated, result) + + +# --------------------------------------------------------------------------- +# Case: per-mb CP/SP reduction (_sp_allreduce_kl_metrics) +# --------------------------------------------------------------------------- + + +def _case_sp_partial_sum(device: torch.device) -> None: + """Catches a re-introduction of the ``v * local_n / total_n`` weighting bug: + each rank's contribution must be its raw partial sum, not a per-rank mean. + Rank-0 numbers chosen so per-rank-mean averaging would give a wrong answer.""" + rank = dist.get_rank() + if rank == 0: + n, ratio_sum, clipfrac_sum, lo, hi = 2, 2.0, 0.0, 0.9, 1.1 + else: + n, ratio_sum, clipfrac_sum, lo, hi = 6, 7.5, 2.0, 0.5, 1.8 + + metrics = _make_metrics( + device, + valid_tokens=n, + ratio_mean=ratio_sum, + pg_clipfrac=clipfrac_sum, + ratio_min=lo, + ratio_max=hi, + ) + metric_ops = {"ratio_min": "min", "ratio_max": "max"} + + mr._sp_allreduce_kl_metrics(metrics, metric_ops, dist.group.WORLD) + + total_n = 2 + 6 + expected_ratio_mean = (2.0 + 7.5) / total_n + expected_clipfrac = (0.0 + 2.0) / total_n + + assert metrics["valid_tokens"] == total_n + got_ratio = metrics["ratio_mean"].item() / metrics["valid_tokens"] + got_clip = metrics["pg_clipfrac"].item() / metrics["valid_tokens"] + assert math.isclose(got_ratio, expected_ratio_mean, rel_tol=1e-12), ( + f"[rank {rank}] ratio_mean: got {got_ratio}, expected {expected_ratio_mean}" + ) + assert math.isclose(got_clip, expected_clipfrac, rel_tol=1e-12), ( + f"[rank {rank}] pg_clipfrac: got {got_clip}, expected {expected_clipfrac}" + ) + assert metrics["ratio_min"].item() == 0.5 + assert metrics["ratio_max"].item() == 1.8 + + +# --------------------------------------------------------------------------- +# Case: cross-mb + cross-DP (_accumulate_is_metrics + _finalize_is_metrics) +# --------------------------------------------------------------------------- + + +# Asymmetric valid_tokens per mb verifies (sum, count) bookkeeping under +# uneven contributions. First two mbs go to rank 0, last two to rank 1. +_DP_MBS = [ + {"valid_tokens": 3, "ratio_mean": 3.6, "pg_clipfrac": 1.0, "ratio_min": 0.7, "ratio_max": 1.5}, + {"valid_tokens": 4, "ratio_mean": 4.0, "pg_clipfrac": 0.0, "ratio_min": 0.9, "ratio_max": 1.2}, + {"valid_tokens": 2, "ratio_mean": 2.5, "pg_clipfrac": 2.0, "ratio_min": 0.4, "ratio_max": 1.8}, + {"valid_tokens": 5, "ratio_mean": 4.5, "pg_clipfrac": 1.0, "ratio_min": 1.1, "ratio_max": 1.3}, +] + + +def _case_dp_accumulate_finalize(device: torch.device) -> None: + rank = dist.get_rank() + my_mbs = _DP_MBS[:2] if rank == 0 else _DP_MBS[2:] + metric_ops = {"ratio_min": "min", "ratio_max": "max"} + + accumulated = {} + for mb in my_mbs: + mr.ModelRunner._accumulate_is_metrics(accumulated, _make_metrics(device, **mb), metric_ops) + + result = {} + _finalize_with_world_dp(accumulated, result) + + total_n = sum(mb["valid_tokens"] for mb in _DP_MBS) + expected = { + "is_ratio_mean": sum(mb["ratio_mean"] for mb in _DP_MBS) / total_n, + "is_pg_clipfrac": sum(mb["pg_clipfrac"] for mb in _DP_MBS) / total_n, + # valid_tokens is itself accumulated as a mean, but its per-mb count + # is +=1 (not +=n_tokens), so the finalized value is total_n / num_mbs. + "is_valid_tokens": total_n / len(_DP_MBS), + "is_ratio_min": min(mb["ratio_min"] for mb in _DP_MBS), + "is_ratio_max": max(mb["ratio_max"] for mb in _DP_MBS), + } + for key, want in expected.items(): + got = result[key] + assert math.isclose(got, want, rel_tol=1e-12), f"[rank {rank}] {key}: got {got}, expected {want}" + + +# --------------------------------------------------------------------------- +# Case: empty-rank min/max identity +# --------------------------------------------------------------------------- + + +def _case_empty_rank_one_empty(device: torch.device) -> None: + """Rank 0 empty (all IGNORE_INDEX → ±inf identity); rank 1 has real values. + The empty rank must not leak into min/max, and the global mean must + reflect rank 1's contribution alone.""" + rank = dist.get_rank() + if rank == 0: + new_metrics = _make_metrics( + device, valid_tokens=0, ratio_mean=0.0, ratio_min=float("inf"), ratio_max=float("-inf") + ) + else: + new_metrics = _make_metrics(device, valid_tokens=5, ratio_mean=6.25, ratio_min=0.6, ratio_max=1.7) + + accumulated = {} + mr.ModelRunner._accumulate_is_metrics(accumulated, new_metrics, {"ratio_min": "min", "ratio_max": "max"}) + result = {} + _finalize_with_world_dp(accumulated, result) + + assert math.isclose(result["is_ratio_mean"], 1.25, rel_tol=1e-12) + assert math.isclose(result["is_ratio_min"], 0.6, rel_tol=1e-12) + assert math.isclose(result["is_ratio_max"], 1.7, rel_tol=1e-12) + + +def _case_empty_rank_all_empty(device: torch.device) -> None: + """All ranks empty. Mean keys with global count == 0 are dropped; min/max + with non-finite reductions fall back to 1.0.""" + new_metrics = _make_metrics(device, valid_tokens=0, ratio_mean=0.0, ratio_min=float("inf"), ratio_max=float("-inf")) + + accumulated = {} + mr.ModelRunner._accumulate_is_metrics(accumulated, new_metrics, {"ratio_min": "min", "ratio_max": "max"}) + result = {} + _finalize_with_world_dp(accumulated, result) + + assert result["is_ratio_min"] == 1.0 + assert result["is_ratio_max"] == 1.0 + assert "is_ratio_mean" not in result, f"empty-rank fallback should drop mean keys: {result}" + + +# --------------------------------------------------------------------------- +# Subprocess dispatch +# --------------------------------------------------------------------------- + + +_CASES = { + "sp_partial_sum": [_case_sp_partial_sum], + "dp_accumulate_finalize": [_case_dp_accumulate_finalize], + "empty_rank": [_case_empty_rank_one_empty, _case_empty_rank_all_empty], +} + + +def _main() -> None: + case_name = os.environ["XORL_TEST_CASE"] + device = _setup_dist() + try: + for fn in _CASES[case_name]: + fn(device) + finally: + dist.destroy_process_group() + + +def _launch(case: str): + return run_distributed_script(__file__, num_gpus=2, timeout=120, extra_env={"XORL_TEST_CASE": case}) + + +if __name__ != "__main__": + + @skip_if_gpu_count_less_than(2) + def test_sp_allreduce_kl_metrics_under_cp(): + _launch("sp_partial_sum").assert_success("CP _sp_allreduce_kl_metrics partial-sum reduction") + + @skip_if_gpu_count_less_than(2) + def test_accumulate_finalize_under_dp(): + _launch("dp_accumulate_finalize").assert_success("DP _accumulate_is_metrics + _finalize_is_metrics") + + @skip_if_gpu_count_less_than(2) + def test_empty_rank_min_max_identity(): + _launch("empty_rank").assert_success("Empty-rank min/max identity fallback") + + +if __name__ == "__main__": + _main() diff --git a/tests/distributed/test_muon_full_gradient.py b/tests/distributed/test_muon_full_gradient.py new file mode 100644 index 00000000..37954a88 --- /dev/null +++ b/tests/distributed/test_muon_full_gradient.py @@ -0,0 +1,208 @@ +"""Distributed correctness tests for Muon ``distributed_mode='full_gradient'``. + +Verifies that the full-gradient path produces the same parameter update as +running Muon on a single rank with the unsharded gradient (i.e., recovers +exact Muon math under FSDP2/DTensor sharding), and that the existing +``shard_local`` mode does NOT match — sanity check that we are exercising +the new code path. +""" + +from __future__ import annotations + +import os +import sys +from pathlib import Path + +import pytest +import torch +import torch.distributed as dist +from torch.distributed.tensor import Shard, distribute_tensor + +from xorl.distributed.parallel_state import get_parallel_state, init_parallel_state +from xorl.optim.muon import Muon +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] + + +def _local_rank() -> int: + return int(os.environ["LOCAL_RANK"]) + + +def _world_size() -> int: + return int(os.environ["WORLD_SIZE"]) + + +def _setup_dist(): + local_rank = _local_rank() + torch.cuda.set_device(local_rank) + dist.init_process_group(backend=get_nccl_backend()) + init_parallel_state( + dp_size=_world_size(), + dp_replicate_size=1, + dp_shard_size=_world_size(), + tp_size=1, + ep_size=1, + pp_size=1, + ulysses_size=1, + ringattn_size=1, + dp_mode="fsdp2", + ) + return torch.device("cuda", local_rank) + + +def _single_rank_oracle(weight_full: torch.Tensor, grad_full: torch.Tensor, *, mode: str) -> torch.Tensor: + """Run Muon on a single rank with the full unsharded gradient and return updated weight.""" + p = torch.nn.Parameter(weight_full.clone()) + p.grad = grad_full.clone() + opt = Muon( + [{"params": [p], "lr": 0.1, "use_muon": True, "weight_decay": 0.0}], + lr=0.1, + momentum=0.0, + nesterov=False, + ns_steps=5, + weight_decay=0.0, + distributed_mode=mode, + ) + opt.step() + return p.detach() + + +def _full_tensor(d): + if hasattr(d, "full_tensor"): + return d.full_tensor() + return d + + +def _layout_shape_and_placements(layout: str): + """Return (global_shape, placements) for a named test layout.""" + if layout == "linear_2d": + # Plain Linear weight, row-sharded on dim 0 — the dense FSDP2 case. + return (16, 12), [Shard(0)] + if layout == "moe_experts_3d": + # 3D MoE expert weight [E, H, I] sharded on dim 1 — mirrors how + # ``parallelize_model_fsdp2`` shards EP-experts (Shard(1) on ep_fsdp). + # Tests the deferred 3D-reshape branch in ``_muon_step``. + return (4, 32, 64), [Shard(1)] + raise ValueError(f"unknown layout: {layout!r}") + + +def _run(distributed_mode: str, layout: str) -> None: + device = _setup_dist() + mesh = get_parallel_state().fsdp_mesh + + torch.manual_seed(42) + global_shape, placements = _layout_shape_and_placements(layout) + weight_full = torch.randn(*global_shape, device=device, dtype=torch.float32) + grad_full = torch.randn(*global_shape, device=device, dtype=torch.float32) + # Make every rank see the same global init by broadcasting from rank 0. + dist.broadcast(weight_full, src=0) + dist.broadcast(grad_full, src=0) + + weight_d = distribute_tensor(weight_full.clone(), mesh, placements) + p = torch.nn.Parameter(weight_d) + + grad_d = distribute_tensor(grad_full.clone(), mesh, placements) + p.grad = grad_d + + opt = Muon( + [{"params": [p], "lr": 0.1, "use_muon": True, "weight_decay": 0.0}], + lr=0.1, + momentum=0.0, + nesterov=False, + ns_steps=5, + weight_decay=0.0, + distributed_mode=distributed_mode, + ) + opt.step() + + full_after = _full_tensor(p.data) + + # Build a single-process oracle running full-gradient mode and compare. + expected_full_grad = _single_rank_oracle(weight_full, grad_full, mode="full_gradient") + + if dist.get_rank() == 0: + if distributed_mode == "full_gradient": + err = (full_after - expected_full_grad).abs().max().item() + assert err < 1e-4, ( + f"[layout={layout}] full_gradient mode does not match single-rank oracle: max abs err = {err}" + ) + print(f"[rank 0] [{layout}] full_gradient max err vs oracle: {err:.6e}") + else: + # shard_local mode: NS runs per local-shard. Should NOT match the + # full-gradient oracle on multi-rank — the orthogonalization is + # genuinely different between shard-local row-slabs (or H-strips + # for 3D) and the full matrix. + err = (full_after - expected_full_grad).abs().max().item() + print(f"[rank 0] [{layout}] shard_local diff vs full-gradient oracle: {err:.6e}") + if dist.get_world_size() > 1: + assert err > 1e-4, ( + f"[layout={layout}] shard_local should differ from full-gradient oracle " + f"on multi-rank; got max abs err = {err} (test premise broken if too small)" + ) + + dist.barrier() + dist.destroy_process_group() + + +def _main() -> None: + mode = os.environ.get("XORL_TEST_MUON_MODE", "full_gradient") + layout = os.environ.get("XORL_TEST_MUON_LAYOUT", "linear_2d") + _run(mode, layout) + + +if __name__ != "__main__": + + @skip_if_gpu_count_less_than(2) + def test_full_gradient_matches_single_rank_oracle_2d(): + result = run_distributed_script( + __file__, + num_gpus=2, + timeout=180, + extra_env={"XORL_TEST_MUON_MODE": "full_gradient", "XORL_TEST_MUON_LAYOUT": "linear_2d"}, + ) + result.assert_success("2D Shard(0) full_gradient Muon should match single-rank oracle") + + @skip_if_gpu_count_less_than(2) + def test_shard_local_differs_from_full_gradient_oracle_2d(): + result = run_distributed_script( + __file__, + num_gpus=2, + timeout=180, + extra_env={"XORL_TEST_MUON_MODE": "shard_local", "XORL_TEST_MUON_LAYOUT": "linear_2d"}, + ) + result.assert_success("2D shard_local should differ from full-gradient oracle on >1 rank") + + @skip_if_gpu_count_less_than(2) + def test_full_gradient_matches_single_rank_oracle_3d_moe(): + # Exercises the deferred-reshape path in ``_muon_step`` for an EP-experts-style + # 3D weight ``[E, H, I]`` sharded on ``H`` (Shard(1)). + result = run_distributed_script( + __file__, + num_gpus=2, + timeout=180, + extra_env={"XORL_TEST_MUON_MODE": "full_gradient", "XORL_TEST_MUON_LAYOUT": "moe_experts_3d"}, + ) + result.assert_success("3D Shard(1) full_gradient Muon should match single-rank oracle") + + @skip_if_gpu_count_less_than(2) + def test_shard_local_differs_from_full_gradient_oracle_3d_moe(): + result = run_distributed_script( + __file__, + num_gpus=2, + timeout=180, + extra_env={"XORL_TEST_MUON_MODE": "shard_local", "XORL_TEST_MUON_LAYOUT": "moe_experts_3d"}, + ) + result.assert_success("3D shard_local should differ from full-gradient oracle on >1 rank") + + +if __name__ == "__main__": + _main() diff --git a/tests/distributed/test_offloading.py b/tests/distributed/test_offloading.py new file mode 100644 index 00000000..84c162cf --- /dev/null +++ b/tests/distributed/test_offloading.py @@ -0,0 +1,19 @@ +import torch + +from xorl.distributed.offloading import build_activation_offloading_context + + +def test_activation_offload_none_gpu_limit_defaults_to_zero() -> None: + model_fwd_context, _ = build_activation_offloading_context( + enable_activation_offload=True, + enable_gradient_checkpointing=False, + activation_gpu_limit=None, + ) + + x = torch.randn(4, 4, requires_grad=True) + with model_fwd_context: + loss = (x * x).sum() + + loss.backward() + + assert x.grad is not None diff --git a/tests/distributed/test_olmo2_qk_rms_norm.py b/tests/distributed/test_olmo2_qk_rms_norm.py new file mode 100644 index 00000000..afd91969 --- /dev/null +++ b/tests/distributed/test_olmo2_qk_rms_norm.py @@ -0,0 +1,112 @@ +"""``Olmo2QKRMSNorm`` regression test. + +OLMo-2's full-axis ``q_norm``/``k_norm`` doesn't compose with stock TP +styles. Under colwise q/k_proj the input arrives hidden-sharded; the +plan wraps these norms with ``LocalAxisRMSNormShard`` to shard their +weight on dim 0. ``Olmo2QKRMSNorm.forward`` detects the Shard(0) +DTensor weight and runs the fused op on local tensors, computing a +local-axis RMS that matches HuggingFace's ``Olmo2RMSNorm`` reference +under TP. + +Can be run two ways: + 1. pytest tests/distributed/test_olmo2_qk_rms_norm.py -v (launches torchrun internally) + 2. torchrun --nproc_per_node=2 tests/distributed/test_olmo2_qk_rms_norm.py (direct) +""" + +import os + +import torch +import torch.distributed as dist +from torch.distributed.device_mesh import init_device_mesh +from torch.distributed.tensor.parallel import parallelize_module + +from xorl.models.layers.normalization import RMSNorm +from xorl.models.transformers.olmo2.modeling_olmo2 import Olmo2QKRMSNorm +from xorl.models.transformers.olmo2.tp_styles import LocalAxisRMSNormShard + + +HIDDEN = 8 +SEQ = 6 +BATCH = 2 + + +def _check_no_tp_passthrough(): + """Without TP, Olmo2QKRMSNorm forward delegates to the parent RMSNorm.""" + norm = Olmo2QKRMSNorm(HIDDEN) + x = torch.randn(BATCH, SEQ, HIDDEN, generator=torch.Generator().manual_seed(0)) + out = norm(x) + assert tuple(out.shape) == (BATCH, SEQ, HIDDEN) + # Numerical equivalence with the parent class on the same input. + ref = RMSNorm(HIDDEN) + ref.weight = torch.nn.Parameter(norm.weight.detach().clone()) + torch.testing.assert_close(out, ref(x), atol=1e-6, rtol=1e-6) + + +def _check_local_axis_rms_norm_shard(mesh): + """LocalAxisRMSNormShard + Olmo2QKRMSNorm: full-axis QK-norm path under TP. + + The colwise q/k_proj output arrives as a plain hidden-sharded tensor; the + custom style shards the weight on dim 0 so each rank's slice matches its + local input. The forward should compute a local-axis RMS matching a + per-rank-local single-tensor ``RMSNorm`` applied to the same shard. + """ + tp = mesh.size() + rank = dist.get_rank() + + norm = parallelize_module(Olmo2QKRMSNorm(HIDDEN), mesh, LocalAxisRMSNormShard()) + + # Mimic colwise q_proj output: same global tensor on every rank, then take + # the rank's hidden slice as the plain (non-DTensor) input. + full = torch.randn(BATCH, SEQ, HIDDEN, generator=torch.Generator().manual_seed(7)) + rank_slice = slice(rank * (HIDDEN // tp), (rank + 1) * (HIDDEN // tp)) + local_input = full[..., rank_slice].contiguous() + + out = norm(local_input) + expected_shape = (BATCH, SEQ, HIDDEN // tp) + assert tuple(out.shape) == expected_shape, f"expected output shape {expected_shape}, got {tuple(out.shape)}" + + # Reference: a plain (non-TP) RMSNorm with hidden=HIDDEN/tp and the local + # weight slice. This is the local-axis RMS HF's OLMo-2 reference computes. + ref_norm = RMSNorm(HIDDEN // tp) + ref_norm.weight = torch.nn.Parameter(norm.weight.to_local().clone()) + ref_out = ref_norm(local_input) + + torch.testing.assert_close(out, ref_out, atol=1e-6, rtol=1e-6) + + +def main(): + dist.init_process_group(backend="gloo") + rank = dist.get_rank() + world_size = dist.get_world_size() + assert HIDDEN % world_size == 0 and SEQ % world_size == 0, ( + "Test fixtures require HIDDEN and SEQ divisible by world_size" + ) + + mesh = init_device_mesh("cpu", (world_size,), mesh_dim_names=("tp",)) + + _check_no_tp_passthrough() + _check_local_axis_rms_norm_shard(mesh) + + if rank == 0: + print("All Olmo2QKRMSNorm checks passed!") + + dist.destroy_process_group() + + +if __name__ != "__main__": + import pytest + + from tests.distributed.distributed_utils import run_distributed_script + + SCRIPT_PATH = os.path.abspath(__file__) + + @pytest.mark.cpu + @pytest.mark.distributed + def test_olmo2_qk_rms_norm_2rank_cpu(): + """Olmo2QKRMSNorm + LocalAxisRMSNormShard on a 2-rank gloo mesh.""" + result = run_distributed_script(SCRIPT_PATH, num_gpus=2, timeout=120) + result.assert_success() + + +if __name__ == "__main__": + main() diff --git a/tests/distributed/test_olmo2_tp_e2e.py b/tests/distributed/test_olmo2_tp_e2e.py new file mode 100644 index 00000000..9e76d12d --- /dev/null +++ b/tests/distributed/test_olmo2_tp_e2e.py @@ -0,0 +1,137 @@ +"""End-to-end OLMo-2 TP=2 forward + backward on a 2-rank gloo+CPU mesh. + +Builds a tiny OLMo-2 model, applies the TP plan from +``xorl.models.transformers.olmo2.parallelize``, runs a full forward, calls +``lm_head`` explicitly, then backprops a scalar loss and checks that every +TP-wrapped parameter received a gradient. Exercises the OLMo-2-specific +machinery the plan introduces: full-axis ``q_norm``/``k_norm`` wrapped +with ``LocalAxisRMSNormShard`` (Shard(0) weight + plain hidden-sharded +input from colwise q/k_proj), driven by ``Olmo2QKRMSNorm.forward`` running +the fused op on local tensors. The rest of the plan is stock +colwise/rowwise (no SP, no PrepareModuleInput). + +Can be run two ways: + 1. pytest tests/distributed/test_olmo2_tp_e2e.py -v (launches torchrun internally) + 2. torchrun --nproc_per_node=2 tests/distributed/test_olmo2_tp_e2e.py (direct) +""" + +import os + +import torch +import torch.distributed as dist +from torch.distributed.device_mesh import init_device_mesh +from torch.distributed.tensor.parallel import parallelize_module + +from xorl.distributed.parallel_state import init_parallel_state +from xorl.distributed.torch_parallelize import _build_tp_plan +from xorl.models.transformers.olmo2.configuration_olmo2 import Olmo2Config +from xorl.models.transformers.olmo2.modeling_olmo2 import Olmo2ForCausalLM + + +def _make_config(): + cfg = Olmo2Config( + architectures=["Olmo2ForCausalLM"], + vocab_size=64, + hidden_size=16, + intermediate_size=32, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=4, + max_position_embeddings=32, + rope_theta=500000.0, + initializer_range=0.5, + attention_dropout=0.0, + tie_word_embeddings=False, + use_cache=False, + ) + cfg._attn_implementation = "eager" + cfg._activation_native = True + return cfg + + +def main(): + dist.init_process_group(backend="gloo") + rank = dist.get_rank() + world_size = dist.get_world_size() + + # OLMo-2's Olmo2Model.forward calls get_parallel_state() to pick a CP + # strategy; initialize it as TP-only so NoopStrategy fires (no all-to-all + # for ulysses, no ring-attention plumbing), keeping the test focused on + # the TP plan we're validating. + init_parallel_state( + dp_size=1, + tp_size=world_size, + dp_mode="none", + device_type="cpu", + ) + mesh = init_device_mesh("cpu", (world_size,), mesh_dim_names=("tp",)) + + torch.manual_seed(0) + model = Olmo2ForCausalLM(_make_config()) + + # Apply unfuse + TP plan via xorl's standard pipeline so we exercise the + # same code path the production launcher uses. + model.unfuse_for_tp() + plan = _build_tp_plan(model) + parallelize_module(model, mesh, plan) + + # Forward. + seq_len = 8 + input_ids = torch.randint(0, 64, (1, seq_len), generator=torch.Generator().manual_seed(rank)) + out = model(input_ids=input_ids) + last_hidden = out.last_hidden_state + + # last_hidden_state comes out of the (un-wrapped) final norm with full + # hidden replicated across ranks — the o_proj/down_proj rowwise default + # all-reduces to Replicate, post-norms are NOT in the TP plan, and the + # residual stream stays Replicate end-to-end. This shape matches what + # vocab_parallel_cross_entropy expects (full [B, S, H] per rank). + assert tuple(last_hidden.shape) == (1, seq_len, 16), ( + f"unexpected last_hidden_state shape: got {tuple(last_hidden.shape)}, expected (1, {seq_len}, 16)" + ) + assert torch.isfinite(last_hidden).all(), "non-finite values in last_hidden_state" + + # Run lm_head explicitly (Olmo2ForCausalLM.forward returns the pre-lm-head + # last_hidden_state; the loss head normally wraps it). lm_head is + # ColwiseParallel — local input + vocab-sharded weight produce a local + # vocab-parallel logits tensor. + logits = model.lm_head(last_hidden) + expected_logits = (1, seq_len, 64 // world_size) + assert tuple(logits.shape) == expected_logits, ( + f"unexpected logits shape: got {tuple(logits.shape)}, expected {expected_logits}" + ) + + # Backward pass: take a scalar loss off the logits and check that gradients + # reach every TP-wrapped parameter without raising. This catches the + # redistribute/grad-flow paths a pure forward would miss + # (LocalAxisRMSNormShard partial reduction, ColwiseParallel/RowwiseParallel + # input/output redistribute backward, lm_head colwise). + loss = logits.float().pow(2).mean() + loss.backward() + no_grad = [name for name, p in model.named_parameters() if p.requires_grad and p.grad is None] + assert not no_grad, f"parameters missing gradients: {no_grad[:5]}" + + if rank == 0: + print("OLMo-2 TP=2 fwd+bwd succeeded; last_hidden_state shape:", tuple(last_hidden.shape)) + + dist.destroy_process_group() + + +if __name__ != "__main__": + import pytest + + from tests.distributed.distributed_utils import run_distributed_script + + SCRIPT_PATH = os.path.abspath(__file__) + + @pytest.mark.cpu + @pytest.mark.distributed + def test_olmo2_tp2_fwd_bwd_e2e_cpu(): + """OLMo-2 fwd+bwd survives the full torchtitan-style TP plan on a 2-rank gloo mesh.""" + result = run_distributed_script(SCRIPT_PATH, num_gpus=2, timeout=180) + result.assert_success() + + +if __name__ == "__main__": + main() diff --git a/tests/distributed/test_sync_padding.py b/tests/distributed/test_sync_padding.py new file mode 100644 index 00000000..a98900c0 --- /dev/null +++ b/tests/distributed/test_sync_padding.py @@ -0,0 +1,71 @@ +from unittest.mock import Mock + +import torch + +from xorl.data.constants import IGNORE_INDEX +from xorl.distributed import sync_padding + + +def test_synchronize_micro_batch_padding_extends_cu_seqlens(monkeypatch): + monkeypatch.setattr(sync_padding.dist, "is_initialized", lambda: True) + monkeypatch.setattr(sync_padding.dist, "get_world_size", lambda group=None: 8) + monkeypatch.setattr(sync_padding, "get_device_type", lambda: "cpu") + + ps = Mock() + ps.cp_enabled = False + monkeypatch.setattr(sync_padding, "get_parallel_state", lambda: ps) + + def fake_all_reduce(tensor, op=None, group=None): + tensor.fill_(512) + + monkeypatch.setattr(sync_padding.dist, "all_reduce", fake_all_reduce) + + micro_batches = [ + { + "input_ids": torch.ones(1, 176, dtype=torch.long), + "labels": torch.ones(1, 176, dtype=torch.long), + "position_ids": torch.arange(176).unsqueeze(0), + "attention_mask": torch.ones(1, 176, dtype=torch.long), + "cu_seq_lens_q": torch.tensor([0, 83, 167, 176], dtype=torch.int32), + "cu_seq_lens_k": torch.tensor([0, 83, 167, 176], dtype=torch.int32), + "max_length_q": torch.tensor(84, dtype=torch.int32), + "max_length_k": torch.tensor(84, dtype=torch.int32), + }, + { + "input_ids": torch.ones(1, 512, dtype=torch.long), + "labels": torch.cat( + [ + torch.ones(176, dtype=torch.long), + torch.full((336,), IGNORE_INDEX, dtype=torch.long), + ] + ).unsqueeze(0), + "position_ids": torch.arange(512).unsqueeze(0), + "attention_mask": torch.cat( + [ + torch.ones(176, dtype=torch.long), + torch.zeros(336, dtype=torch.long), + ] + ).unsqueeze(0), + "cu_seq_lens_q": torch.tensor([0, 83, 167, 176], dtype=torch.int32), + "cu_seq_lens_k": torch.tensor([0, 83, 167, 176], dtype=torch.int32), + "max_length_q": torch.tensor(84, dtype=torch.int32), + "max_length_k": torch.tensor(84, dtype=torch.int32), + }, + ] + + sync_padding.synchronize_micro_batch_padding(micro_batches) + + for mb in micro_batches: + assert mb["input_ids"].shape[-1] == 512 + assert mb["labels"].shape[-1] == 512 + assert mb["position_ids"].shape[-1] == 512 + assert mb["attention_mask"].shape[-1] == 512 + assert torch.equal(mb["labels"][0, 176:], torch.full((336,), IGNORE_INDEX)) + assert mb["attention_mask"][0, 176:].sum().item() == 0 + + assert mb["cu_seq_lens_q"].tolist() == [0, 83, 167, 512] + assert mb["cu_seq_lens_k"].tolist() == [0, 83, 167, 512] + assert mb["max_length_q"] == 345 + assert mb["max_length_k"] == 345 + assert isinstance(mb["max_length_q"], int) + assert isinstance(mb["max_length_k"], int) diff --git a/tests/distributed/test_vocab_parallel_ce.py b/tests/distributed/test_vocab_parallel_ce.py index 37fe23fa..5203ceea 100644 --- a/tests/distributed/test_vocab_parallel_ce.py +++ b/tests/distributed/test_vocab_parallel_ce.py @@ -251,22 +251,6 @@ def test_vocab_parallel_ce_2gpu(): result = run_distributed_script(SCRIPT_PATH, num_gpus=2, timeout=180) result.assert_success() - @pytest.mark.gpu - @pytest.mark.distributed - @skip_if_gpu_count_less_than(4) - def test_vocab_parallel_ce_4gpu(): - """Vocab-parallel cross-entropy correctness + backward with 4 GPUs.""" - result = run_distributed_script(SCRIPT_PATH, num_gpus=4, timeout=180) - result.assert_success() - - @pytest.mark.gpu - @pytest.mark.distributed - @skip_if_gpu_count_less_than(8) - def test_vocab_parallel_ce_8gpu(): - """Vocab-parallel cross-entropy correctness + backward with 8 GPUs.""" - result = run_distributed_script(SCRIPT_PATH, num_gpus=8, timeout=180) - result.assert_success() - if __name__ == "__main__": main() diff --git a/tests/e2e/e2e_utils.py b/tests/e2e/e2e_utils.py index d45f7289..8dc5021a 100644 --- a/tests/e2e/e2e_utils.py +++ b/tests/e2e/e2e_utils.py @@ -12,7 +12,6 @@ import math import os import random -import socket import subprocess import sys from dataclasses import dataclass @@ -26,6 +25,8 @@ from tokenizers.pre_tokenizers import Whitespace from transformers import AutoConfig, AutoModelForCausalLM +from .server_utils import _get_free_port + try: import xorl_client @@ -418,17 +419,6 @@ def generate_training_config( return yaml_path -# --------------------------------------------------------------------------- -# Training launcher -# --------------------------------------------------------------------------- - - -def _get_free_port() -> int: - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: - s.bind(("", 0)) - return s.getsockname()[1] - - def run_training( config_path: str, num_gpus: int = 1, diff --git a/tests/e2e/qwen3_8b/test_distsignsgd.py b/tests/e2e/qwen3_8b/test_distsignsgd.py new file mode 100644 index 00000000..3814fa4e --- /dev/null +++ b/tests/e2e/qwen3_8b/test_distsignsgd.py @@ -0,0 +1,71 @@ +"""E2E tests for DistSignSGD with tiny Qwen3 dense models.""" + +import math +import os + +import pytest + +from tests.e2e.e2e_utils import ( + generate_training_config, + run_training, + skip_if_gpu_count_less_than, +) + + +pytestmark = [pytest.mark.e2e, pytest.mark.gpu, pytest.mark.slow] + + +class TestDistSignSGD2GPU: + @skip_if_gpu_count_less_than(2) + def test_distsignsgd_fsdp2_runs_with_finite_loss(self, tiny_dense_model_dir): + """Trainer: DistSignSGD on 2-GPU FSDP2 completes and reports finite metrics.""" + output_dir = os.path.join(tiny_dense_model_dir, "output_distsignsgd_fsdp2") + config_path = generate_training_config( + model_dir=tiny_dense_model_dir, + output_dir=output_dir, + num_gpus=2, + dp_shard_size=2, + gradient_accumulation_steps=2, + packing_seq_len=256, + optimizer="distsignsgd", + max_steps=5, + lr=1e-3, + ) + + result = run_training(config_path, num_gpus=2, timeout=600) + + result.assert_success() + assert result.global_step == 5 + assert result.final_loss is not None and not math.isnan(result.final_loss) + assert result.final_loss < 12.0 + assert result.loss_history is not None and len(result.loss_history) == 5 + + @skip_if_gpu_count_less_than(4) + def test_distsignsgd_with_ulysses_outside_fsdp_and_gradient_accumulation_runs( + self, + tiny_dense_model_dir_with_weights, + ): + """Trainer: DistSignSGD keeps Ulysses exact-sum outside FSDP and accumulates mean sign votes.""" + output_dir = os.path.join(tiny_dense_model_dir_with_weights, "output_distsignsgd_fsdp2_u2_dp2_gas2") + config_path = generate_training_config( + model_dir=tiny_dense_model_dir_with_weights, + model_path=tiny_dense_model_dir_with_weights, + output_dir=output_dir, + num_gpus=4, + dp_shard_size=2, + ulysses_size=2, + gradient_accumulation_steps=2, + packing_seq_len=256, + optimizer="distsignsgd", + max_steps=5, + lr=1e-3, + extra_train={"cp_fsdp_mode": "none"}, + ) + + result = run_training(config_path, num_gpus=4, timeout=900) + + result.assert_success() + assert result.global_step == 5 + assert result.final_loss is not None and not math.isnan(result.final_loss) + assert result.final_loss < 12.0 + assert result.loss_history is not None and len(result.loss_history) == 5 diff --git a/tests/e2e/qwen3_8b/test_tflops_threshold.py b/tests/e2e/qwen3_8b/test_tflops_threshold.py index 4a03286f..fbb68611 100644 --- a/tests/e2e/qwen3_8b/test_tflops_threshold.py +++ b/tests/e2e/qwen3_8b/test_tflops_threshold.py @@ -100,12 +100,11 @@ def _run_tflops_benchmark(num_gpus: int, dp_shard_size: int, tmp_path: str): _, 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) - adam_params = xorl_client.AdamParams(learning_rate=1e-4, beta1=0.9, beta2=0.95, eps=1e-8) # Warmup for _ in range(NUM_WARMUP): training_client.forward_backward(data, loss_fn="causallm_loss").result() - training_client.optim_step(adam_params).result() + training_client.optim_step(learning_rate=1e-4).result() # Measured steps step_times = [] @@ -113,7 +112,7 @@ def _run_tflops_benchmark(num_gpus: int, dp_shard_size: int, tmp_path: str): for step in range(NUM_STEPS): t0 = time.perf_counter() fwd_bwd = training_client.forward_backward(data, loss_fn="causallm_loss") - optim = training_client.optim_step(adam_params) + optim = training_client.optim_step(learning_rate=1e-4) result = fwd_bwd.result() optim.result() t1 = time.perf_counter() diff --git a/tests/e2e/server_utils.py b/tests/e2e/server_utils.py index 7d57ebb6..b4aac3d2 100644 --- a/tests/e2e/server_utils.py +++ b/tests/e2e/server_utils.py @@ -8,7 +8,8 @@ import subprocess import sys import time -from typing import List, Optional +from pathlib import Path +from typing import Any, List, Optional import pytest import requests @@ -190,6 +191,53 @@ def _wait_for_server(url: str, timeout: float = 120.0, poll_interval: float = 2. return False +def wait_for_future(base_url: str, request_id: str, timeout: float = 240.0) -> dict: + deadline = time.time() + timeout + while time.time() < deadline: + response = requests.post(f"{base_url}/api/v1/retrieve_future", json={"request_id": request_id}, timeout=60) + response.raise_for_status() + payload = response.json() + if payload.get("type") == "try_again": + time.sleep(0.5) + continue + if payload.get("type") == "request_failed" or payload.get("error"): + raise AssertionError(f"Future {request_id} failed: {payload}") + return payload + raise TimeoutError(f"Future {request_id} timed out after {timeout}s") + + +def post_and_wait_for_future( + base_url: str, + endpoint: str, + payload: dict[str, Any], + *, + request_timeout: float = 30.0, + future_timeout: float = 240.0, +) -> dict: + response = requests.post(f"{base_url}{endpoint}", json=payload, timeout=request_timeout) + response.raise_for_status() + return wait_for_future(base_url, response.json()["request_id"], timeout=future_timeout) + + +def scalar(value: Any) -> float: + if isinstance(value, dict) and "data" in value: + return float(value["data"][0]) + if hasattr(value, "data"): + return float(value.data[0]) + return float(value) + + +def get_sglang_paths() -> tuple[Path | None, Path | None]: + internal_dir = os.environ.get("XORL_SGLANG_INTERNAL_DIR") + python_path = os.environ.get("XORL_SGLANG_PYTHON") + source_dir = os.environ.get("XORL_SGLANG_SOURCE_DIR") + if internal_dir: + root = Path(internal_dir) + python_path = python_path or str(root / ".venv" / "bin" / "python") + source_dir = source_dir or str(root / "python") + return (Path(python_path) if python_path else None, Path(source_dir) if source_dir else None) + + _PORT_RETRY_ATTEMPTS = 3 @@ -369,14 +417,12 @@ def extract_loss(fwd_bwd_result) -> float: def run_sft_steps(training_client, data, num_steps=5, lr=1e-3) -> list: """Run SFT training steps and return loss history.""" - xorl_client = _require_xorl_client() - - adam_params = xorl_client.AdamParams(learning_rate=lr, beta1=0.9, beta2=0.95, eps=1e-8) + _require_xorl_client() losses = [] for step in range(num_steps): fwd_bwd = training_client.forward_backward(data, loss_fn="causallm_loss") - optim = training_client.optim_step(adam_params) + optim = training_client.optim_step(learning_rate=lr) result = fwd_bwd.result() optim.result() diff --git a/tests/e2e/test_opd_cpu.py b/tests/e2e/test_opd_cpu.py new file mode 100644 index 00000000..611d9090 --- /dev/null +++ b/tests/e2e/test_opd_cpu.py @@ -0,0 +1,246 @@ +"""CPU e2e-style validation for OPD server data path.""" + +import asyncio + +import pytest +import torch + +from tests._helpers.opd import make_teacher_files, reference_grouped_opd_loss +from xorl.distillation import TeacherActivationCache, TeacherHeadManager +from xorl.ops.loss import opd_loss_function +from xorl.server.backend import DummyBackend +from xorl.server.orchestrator.request_processor import RequestProcessor +from xorl.server.protocol.api_orchestrator import OrchestratorRequest, OutputType +from xorl.server.protocol.operations import ModelPassData +from xorl.server.runner.utils.batch_utils import convert_batch_to_tensors + + +pytestmark = [pytest.mark.e2e, pytest.mark.cpu, pytest.mark.server] + + +class OPDCPUBackend(DummyBackend): + def __init__(self, student_hidden_table: torch.Tensor, student_head: torch.Tensor): + super().__init__() + self.student_hidden_table = student_hidden_table + self.student_head = student_head + self.last_batches = None + self.last_routed_experts = None + self.last_routed_expert_logits = None + + async def forward_backward( + self, + batches, + loss_fn="causallm_loss", + loss_fn_params=None, + model_id=None, + routed_experts=None, + routed_expert_logits=None, + request_id=None, + ): + assert loss_fn == "opd_loss" + assert model_id == "opd-e2e" + self.last_batches = batches + self.last_routed_experts = routed_experts + self.last_routed_expert_logits = routed_expert_logits + + params = loss_fn_params or {} + head_manager = TeacherHeadManager(params["teacher_heads"]) + hidden_cache = TeacherActivationCache(params["teacher_hidden_caches"]) + + total_loss = torch.tensor(0.0) + valid_tokens = 0 + for raw_batch in batches: + batch = convert_batch_to_tensors(raw_batch) + input_ids = batch["input_ids"] + labels = batch["labels"] + teacher_ids = batch["teacher_ids"] + local_valid = int((labels != -100).sum().item()) + valid_tokens += local_valid + + hidden_states = self.student_hidden_table[input_ids] + local_loss = hidden_states.sum() * 0.0 + self.student_head.sum() * 0.0 + for teacher_id in torch.unique(teacher_ids[labels != -100]).tolist(): + teacher_id = int(teacher_id) + teacher_labels = labels.masked_fill(teacher_ids != teacher_id, -100) + teacher_hidden_states = hidden_cache.get( + teacher_id, + batch["teacher_cache_indices"], + device="cpu", + dtype=hidden_states.dtype, + ) + teacher_head = head_manager.get(teacher_id, device="cpu", dtype=self.student_head.dtype) + result = opd_loss_function( + hidden_states=hidden_states, + weight=self.student_head, + labels=teacher_labels, + teacher_hidden_states=teacher_hidden_states, + teacher_lm_head_weight=teacher_head, + teacher_weights=batch["teacher_weights"], + num_chunks=2, + normalization_denominator=torch.tensor(local_valid), + ) + local_loss = local_loss + result.loss + total_loss = total_loss + local_loss + + return { + "total_loss": float(total_loss.item()), + "global_valid_tokens": valid_tokens, + "opd_kl": 0.123, + "opd_num_teachers:max": 2, + "execution_time": 0.0, + } + + +def test_opd_request_processor_to_backend_e2e(tmp_path): + torch.manual_seed(123) + vocab_size = 17 + hidden_size = 6 + teacher_cache_size = 16 + + student_hidden_table = torch.randn(vocab_size, hidden_size) / hidden_size**0.5 + student_head = torch.randn(vocab_size, hidden_size) / hidden_size**0.5 + teacher_heads = { + "0": torch.randn(vocab_size, hidden_size) / hidden_size**0.5, + "1": torch.randn(vocab_size, hidden_size) / hidden_size**0.5, + } + teacher_hidden_caches = { + "0": torch.randn(teacher_cache_size, hidden_size) / hidden_size**0.5, + "1": torch.randn(teacher_cache_size, hidden_size) / hidden_size**0.5, + } + + teacher_files = make_teacher_files(tmp_path, teacher_heads, teacher_hidden_caches) + + backend = OPDCPUBackend(student_hidden_table=student_hidden_table, student_head=student_head) + processor = RequestProcessor( + backend=backend, + sample_packing_sequence_len=16, + enable_packing=True, + pad_to_multiple_of=1, + cp_size=1, + ) + + # Deliberately unsorted by teacher. RequestProcessor should group by teacher + # before packing so a training mini-batch only loads one head at a time. + data = [ + { + "input_ids": [8, 9], + "target_tokens": [9, 10], + "teacher_id": 1, + "teacher_weight": 0.5, + "teacher_cache_indices": [4, 5], + }, + { + "input_ids": [1, 2, 3], + "target_tokens": [2, 3, 4], + "teacher_id": 0, + "teacher_weight": 1.0, + "teacher_cache_indices": [0, 1, 2], + }, + { + "input_ids": [11], + "target_tokens": [12], + "teacher_id": 1, + "teacher_weight": 1.5, + "teacher_cache_indices": [6], + }, + ] + request = OrchestratorRequest( + operation="forward_backward", + payload=ModelPassData( + data=data, + loss_fn="opd_loss", + loss_fn_params={ + "teacher_heads": teacher_files.heads, + "teacher_hidden_caches": teacher_files.hidden_caches, + "opd_sort_by_teacher": True, + }, + model_id="opd-e2e", + routed_experts=["r1-a", "r0", "r1-b"], + routed_expert_logits=["l1-a", "l0", "l1-b"], + ), + ) + + output = asyncio.run(processor.execute_forward_backward(request)) + + assert output.output_type == OutputType.FORWARD_BACKWARD + assert output.finished is True + result = output.outputs[0] + assert result["success"] is True + assert result["valid_tokens"] == 6 + assert result["opd_kl"] == 0.123 + assert result["opd_num_teachers:max"] == 2 + assert "is_opd_kl" not in result + + assert backend.last_routed_experts == ["r0", "r1-a", "r1-b"] + assert backend.last_routed_expert_logits == ["l0", "l1-a", "l1-b"] + assert backend.last_batches[0]["teacher_ids"] == [[0, 0, 0, 1, 1, 1]] + + expected = reference_grouped_opd_loss( + backend.last_batches[0], + student_hidden_table, + student_head, + teacher_hidden_caches, + teacher_heads, + ) + assert result["loss"] == pytest.approx(expected.item(), rel=1e-5, abs=1e-6) + + +class TeacherCacheCPUBackend(DummyBackend): + async def forward( + self, + batches, + loss_fn="causallm_loss", + loss_fn_params=None, + model_id=None, + routed_experts=None, + routed_expert_logits=None, + request_id=None, + ): + assert loss_fn == "teacher_hidden_cache" + assert model_id == "teacher" + return { + "total_loss": 0.0, + "global_valid_tokens": 5, + "teacher_hidden_cache": { + "path": "/tmp/teacher_hidden.safetensors", + "tensor_key": "hidden_states", + "num_tokens": 5, + "hidden_size": 6, + "cache_indices_by_sample": [[0, 1, 2], [3, 4]], + }, + "teacher_prefill_tokens": 5, + "teacher_prefill_forward_compute_s": 0.25, + "teacher_hidden_cache_write_s": 0.01, + "execution_time": 0.3, + } + + +def test_teacher_hidden_cache_metadata_passes_through_request_processor(): + backend = TeacherCacheCPUBackend() + processor = RequestProcessor( + backend=backend, + sample_packing_sequence_len=16, + enable_packing=True, + pad_to_multiple_of=1, + cp_size=1, + ) + request = OrchestratorRequest( + operation="forward", + payload=ModelPassData( + data=[ + {"input_ids": [1, 2, 3], "target_tokens": [1, 2, 3]}, + {"input_ids": [4, 5], "target_tokens": [4, 5]}, + ], + loss_fn="teacher_hidden_cache", + loss_fn_params={"teacher_hidden_cache_path": "/tmp/teacher_hidden.safetensors"}, + model_id="teacher", + ), + ) + + output = asyncio.run(processor.execute_forward(request)) + + assert output.output_type == OutputType.FORWARD + result = output.outputs[0] + assert result["teacher_hidden_cache"]["cache_indices_by_sample"] == [[0, 1, 2], [3, 4]] + assert result["teacher_prefill_tokens"] == 5 + assert result["teacher_prefill_forward_compute_s"] == 0.25 diff --git a/tests/e2e/test_opd_full_pipeline.py b/tests/e2e/test_opd_full_pipeline.py new file mode 100644 index 00000000..cc8c121f --- /dev/null +++ b/tests/e2e/test_opd_full_pipeline.py @@ -0,0 +1,450 @@ +"""Full OPD pipeline e2e test: separate student SGLang sampler + teacher +SGLang hidden-state server + xorl GPU trainer with weight sync. + +This is the closest end-to-end exercise of section 5.1.2 we can run +locally with tiny models. Each iteration of the loop: + + student SGLang --(rollouts)--> teacher SGLang --(hidden_states)--> xorl trainer + ^ | + | v + +----------------- sync_inference_weights (NCCL) ------------------------+ + +It validates the wiring rather than learning quality: the model is random +and the prompts are short, so we only check that every step succeeds, all +losses are finite, the student SGLang weight_version advances after each +sync, and rollouts after sync differ from the initial model output. +""" + +from __future__ import annotations + +import json +import os +import signal +import subprocess +import time +from pathlib import Path + +import pytest +import requests +import torch + +from tests._helpers.opd import save_teacher_hidden_cache + +from .e2e_utils import create_tiny_model_dir +from .server_utils import ( + ServerProcess, + _get_free_port, + _start_server_or_fail, + generate_server_config, + get_sglang_paths, + post_and_wait_for_future, + scalar, +) + + +SGLANG_PYTHON, SGLANG_SOURCE_DIR = get_sglang_paths() + + +pytestmark = [ + pytest.mark.e2e, + pytest.mark.gpu, + pytest.mark.server, + pytest.mark.slow, + pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"), + pytest.mark.skipif(torch.cuda.device_count() < 3, reason="Need >=3 GPUs (trainer + student + teacher)"), + pytest.mark.skipif( + SGLANG_PYTHON is None or not SGLANG_PYTHON.exists(), + reason="Set XORL_SGLANG_PYTHON or XORL_SGLANG_INTERNAL_DIR", + ), + pytest.mark.skipif( + SGLANG_SOURCE_DIR is None or not SGLANG_SOURCE_DIR.exists(), + reason="Set XORL_SGLANG_SOURCE_DIR or XORL_SGLANG_INTERNAL_DIR", + ), +] + + +class _SGLangServer: + """Subprocess wrapper around `python -m sglang.launch_server`. + + Uses a separate CUDA_VISIBLE_DEVICES per server so the student and + teacher are isolated from the xorl trainer's GPU. + """ + + def __init__( + self, + model_dir: str, + gpu_index: int, + log_path: Path, + *, + port: int | None = None, + enable_return_hidden_states: bool = False, + ) -> None: + self.model_dir = model_dir + self.gpu_index = gpu_index + self.log_path = log_path + self.port = port or _get_free_port() + self.enable_return_hidden_states = enable_return_hidden_states + self.process: subprocess.Popen | None = None + self.base_url = f"http://127.0.0.1:{self.port}" + + def start(self, timeout: float = 180.0) -> None: + assert SGLANG_PYTHON is not None + assert SGLANG_SOURCE_DIR is not None + cmd = [ + str(SGLANG_PYTHON), + "-m", + "sglang.launch_server", + "--model-path", + self.model_dir, + "--tokenizer-path", + self.model_dir, + "--host", + "127.0.0.1", + "--port", + str(self.port), + "--dtype", + "bfloat16", + "--attention-backend", + "triton", + "--disable-cuda-graph", + "--disable-piecewise-cuda-graph", + "--disable-radix-cache", + "--mem-fraction-static", + "0.2", + "--max-total-tokens", + "256", + "--max-running-requests", + "4", + "--log-level", + "warning", + # We always pass input_ids in this test, so the tokenizer is never + # actually exercised. Skip warmup (which would otherwise fail trying + # to tokenize "The capital city of France is" with the tiny + # WordLevel vocab) but leave tokenizer init enabled. Note: combining + # --skip-tokenizer-init with weight sync is broken upstream — the + # scheduler sends WeightUpdatePauseReq through the tokenizer pipe in + # that mode and the tokenizer manager has no handler for it. + "--skip-server-warmup", + ] + if self.enable_return_hidden_states: + cmd.append("--enable-return-hidden-states") + + env = os.environ.copy() + env["PYTHONPATH"] = f"{SGLANG_SOURCE_DIR}:{env.get('PYTHONPATH', '')}" + env["CUDA_VISIBLE_DEVICES"] = str(self.gpu_index) + env.pop("CUDA_DEVICE_ORDER", None) + + log_file = self.log_path.open("w") + self._log_file = log_file + self.process = subprocess.Popen( + cmd, + env=env, + stdout=log_file, + stderr=subprocess.STDOUT, + preexec_fn=os.setsid, + ) + + deadline = time.time() + timeout + while time.time() < deadline: + if self.process.poll() is not None: + tail = self._tail() + raise AssertionError( + f"SGLang server (gpu={self.gpu_index}, port={self.port}) exited early.\n--- log tail ---\n{tail}" + ) + try: + resp = requests.get(f"{self.base_url}/health", timeout=3) + if resp.status_code == 200: + return + except requests.exceptions.RequestException: + pass + time.sleep(2.0) + tail = self._tail() + raise TimeoutError( + f"SGLang server (gpu={self.gpu_index}, port={self.port}) not healthy after {timeout}s.\n" + f"--- log tail ---\n{tail}" + ) + + def stop(self) -> None: + if self.process is not None: + try: + os.killpg(os.getpgid(self.process.pid), signal.SIGTERM) + self.process.wait(timeout=20) + except (ProcessLookupError, subprocess.TimeoutExpired): + try: + os.killpg(os.getpgid(self.process.pid), signal.SIGKILL) + self.process.wait(timeout=10) + except Exception: + pass + self.process = None + if hasattr(self, "_log_file") and self._log_file is not None: + self._log_file.close() + self._log_file = None + + def _tail(self, n: int = 60) -> str: + if not self.log_path.exists(): + return "" + with self.log_path.open() as f: + lines = f.readlines() + return "".join(lines[-n:]) + + def model_info(self) -> dict: + resp = requests.get(f"{self.base_url}/model_info", timeout=10) + resp.raise_for_status() + return resp.json() + + def generate( + self, + input_ids: list[int], + *, + max_new_tokens: int, + temperature: float = 0.0, + return_hidden_states: bool = False, + timeout: float = 60.0, + ) -> dict: + payload = { + "input_ids": input_ids, + "sampling_params": {"temperature": temperature, "max_new_tokens": max_new_tokens}, + "return_hidden_states": return_hidden_states, + } + resp = requests.post(f"{self.base_url}/generate", json=payload, timeout=timeout) + resp.raise_for_status() + return resp.json() + + +def _student_sample(student: _SGLangServer, prompt_ids: list[int], max_new_tokens: int) -> list[int]: + """Generate a continuation; return prompt+completion as a flat token list.""" + out = student.generate(prompt_ids, max_new_tokens=max_new_tokens, temperature=0.7) + completion_ids = out.get("output_ids") or out.get("token_ids") or [] + if not completion_ids: + # Some SGLang versions place ids under meta_info instead. + completion_ids = out.get("meta_info", {}).get("output_ids", []) + if isinstance(completion_ids, dict): + completion_ids = completion_ids.get("output_ids", []) + completion_ids = [int(t) for t in completion_ids] + return list(prompt_ids) + completion_ids + + +def _teacher_hidden_states_for(teacher: _SGLangServer, sequences: list[list[int]]) -> list[torch.Tensor]: + """Return one [seq_len, hidden] tensor per input sequence using teacher's prefill hidden states.""" + hiddens: list[torch.Tensor] = [] + for seq in sequences: + out = teacher.generate( + seq, + max_new_tokens=1, + temperature=0.0, + return_hidden_states=True, + ) + chunks = out.get("meta_info", {}).get("hidden_states") or [] + if not chunks: + raise AssertionError(f"Teacher SGLang did not return hidden_states. Response keys: {list(out.keys())}") + prefill = torch.tensor(chunks[0], dtype=torch.bfloat16) + if prefill.ndim != 2 or prefill.shape[0] != len(seq): + raise AssertionError( + f"Unexpected teacher hidden state shape {tuple(prefill.shape)} for input length {len(seq)}" + ) + hiddens.append(prefill) + return hiddens + + +def test_opd_full_pipeline_with_weight_sync(tmp_path): + torch.manual_seed(2026) + + # Models — both student and teacher use the same tiny config so OPD vocab/hidden line up. + student_dir = create_tiny_model_dir(str(tmp_path / "student"), model_type="dense", save_weights=True) + teacher_dir = create_tiny_model_dir(str(tmp_path / "teacher"), model_type="dense", save_weights=True) + + visible = os.environ.get("CUDA_VISIBLE_DEVICES") + visible_devices = ( + [int(x) for x in visible.split(",") if x.strip() != ""] if visible else list(range(torch.cuda.device_count())) + ) + if len(visible_devices) < 3: + pytest.skip("Need at least 3 visible CUDA devices") + trainer_gpu, student_gpu, teacher_gpu = visible_devices[0], visible_devices[1], visible_devices[2] + + teacher_log = tmp_path / "teacher_sglang.log" + student_log = tmp_path / "student_sglang.log" + teacher = _SGLangServer( + model_dir=teacher_dir, + gpu_index=teacher_gpu, + log_path=teacher_log, + enable_return_hidden_states=True, + ) + student = _SGLangServer( + model_dir=student_dir, + gpu_index=student_gpu, + log_path=student_log, + enable_return_hidden_states=False, + ) + + output_dir = tmp_path / "xorl_server" + config_path = generate_server_config( + student_dir, + str(output_dir), + num_gpus=1, + enable_lora=False, + enable_gradient_checkpointing=False, + sample_packing_sequence_len=64, + extra_config={ + "enable_full_shard": False, + "worker_connection_timeout": 240.0, + }, + ) + server = ServerProcess( + config_path=config_path, + num_gpus=1, + api_port=_get_free_port(), + output_dir=str(output_dir), + ) + + # The xorl ServerProcess inherits CUDA_VISIBLE_DEVICES from the parent + # only when set; otherwise it sets it to range(num_gpus). We force the + # trainer onto a single GPU that does not collide with student/teacher. + prev_visible = os.environ.get("CUDA_VISIBLE_DEVICES") + os.environ["CUDA_VISIBLE_DEVICES"] = str(trainer_gpu) + + # TeacherHeadManager understands HF directory paths and pulls lm_head.weight + # straight from the safetensors shard, so no extra extraction step is needed. + teacher_head_entry = teacher_dir + + try: + teacher.start(timeout=300.0) + student.start(timeout=300.0) + try: + _start_server_or_fail(server, timeout=300.0) + finally: + if prev_visible is None: + os.environ.pop("CUDA_VISIBLE_DEVICES", None) + else: + os.environ["CUDA_VISIBLE_DEVICES"] = prev_visible + + train_url = server.base_url + + # Two short OPD steps: rollout -> teacher hidden -> forward_backward + # -> optim_step -> save trained weights to disk -> have student SGLang + # reload from disk. We use the disk-based refresh path (xorl + # /save_weights_for_sampler + SGLang /update_weights_from_disk) here. + # The NCCL `sync_inference_weights` path was wired up and gets all the + # way through the broadcast (5 buckets, ~26 params transferred), but + # SGLang's /destroy_weights_update_group never returns for our tiny + # model setup, which strands the call. The xorl side of that path is + # unblocked by the wait_for_workers=False fix that ships in this branch. + prompts = [[5, 6, 7], [11, 12], [21, 22, 23, 24]] + max_new_tokens = 4 + loss_history: list[float] = [] + rollout_history: list[list[list[int]]] = [] + + for step in range(2): + sequences = [_student_sample(student, p, max_new_tokens) for p in prompts] + rollout_history.append([list(seq) for seq in sequences]) + assert all(len(seq) >= len(p) + 1 for seq, p in zip(sequences, prompts)), ( + f"step {step}: student SGLang produced no completion tokens: {sequences}" + ) + + hiddens = _teacher_hidden_states_for(teacher, sequences) + cache_path = tmp_path / f"teacher_hidden_step{step}.safetensors" + cache_indices = save_teacher_hidden_cache(hiddens, cache_path) + + data = [] + for seq, indices in zip(sequences, cache_indices): + data.append( + { + "model_input": {"input_ids": seq}, + "loss_fn_inputs": { + "target_tokens": seq, + "teacher_ids": [0] * len(seq), + "teacher_weights": [1.0] * len(seq), + "teacher_cache_indices": indices, + }, + } + ) + + fb = post_and_wait_for_future( + train_url, + "/api/v1/forward_backward", + { + "model_id": "default", + "forward_backward_input": { + "data": data, + "loss_fn": "opd_loss", + "loss_fn_params": { + "teacher_heads": {"0": str(teacher_head_entry)}, + "teacher_hidden_caches": {"0": str(cache_path)}, + "opd_sort_by_teacher": True, + "num_chunks": 4, + }, + }, + }, + request_timeout=60.0, + future_timeout=300.0, + ) + loss_val = scalar(fb["loss_fn_outputs"][0]["loss"]) + assert loss_val == loss_val and loss_val >= 0.0, ( # NaN-safe + f"step {step}: loss is invalid ({loss_val})" + ) + loss_history.append(loss_val) + + opt = post_and_wait_for_future( + train_url, + "/api/v1/optim_step", + { + "model_id": "default", + "adam_params": {"learning_rate": 1e-3, "beta1": 0.9, "beta2": 0.95, "eps": 1e-8}, + "gradient_clip": 1.0, + }, + request_timeout=60.0, + future_timeout=120.0, + ) + grad_norm = scalar(opt["metrics"]["grad_norm"]) + assert grad_norm == grad_norm, f"step {step}: NaN grad norm" + + # Persist the trained weights as an HF safetensors directory and + # tell the student SGLang server to reload them in place. This + # refreshes the on-policy sampler with the newest student weights. + sampler_name = f"opd-step-{step}" + save_future = post_and_wait_for_future( + train_url, + "/api/v1/save_weights_for_sampler", + {"model_id": "default", "name": sampler_name}, + request_timeout=60.0, + future_timeout=300.0, + ) + sampler_dir = Path(output_dir) / "sampler_weights" / sampler_name + assert sampler_dir.exists(), ( + f"step {step}: expected sampler weights dir at {sampler_dir} (future={save_future})" + ) + + for fname in ("tokenizer.json", "tokenizer_config.json", "config.json"): + src = Path(student_dir) / fname + dst = sampler_dir / fname + if src.exists() and not dst.exists(): + dst.write_bytes(src.read_bytes()) + + update_resp = requests.post( + f"{student.base_url}/update_weights_from_disk", + json={"model_path": str(sampler_dir), "abort_all_requests": True, "flush_cache": True}, + timeout=300, + ) + update_resp.raise_for_status() + assert update_resp.json().get("success"), ( + f"step {step}: SGLang update_weights_from_disk failed: {update_resp.json()}" + ) + + # Sanity: the model is random and tiny but two consecutive OPD updates with + # weight refresh should produce at least one differing rollout across runs. + assert any(rollout_history[0][i] != rollout_history[1][i] for i in range(len(prompts))), ( + "Rollouts before and after weight refresh are byte-identical; " + "weight refresh may not have taken effect.\n" + f"step0: {rollout_history[0]}\nstep1: {rollout_history[1]}" + ) + + # Persist the loss history for debugging — useful if a CI failure + # only shows the top-level assertion. + (tmp_path / "loss_history.json").write_text(json.dumps(loss_history), encoding="utf-8") + finally: + try: + server.stop() + finally: + try: + student.stop() + finally: + teacher.stop() diff --git a/tests/e2e/test_opd_gpu.py b/tests/e2e/test_opd_gpu.py new file mode 100644 index 00000000..60804c93 --- /dev/null +++ b/tests/e2e/test_opd_gpu.py @@ -0,0 +1,78 @@ +"""CUDA smoke coverage for OPD runner teacher grouping.""" + +import pytest +import torch + +from tests._helpers.opd import make_teacher_files +from xorl.server.runner.model_runner import ModelRunner + + +pytestmark = [ + pytest.mark.e2e, + pytest.mark.gpu, + pytest.mark.server, + pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"), +] + + +def test_opd_runner_grouped_teachers_cuda(tmp_path): + torch.manual_seed(321) + device = torch.device("cuda") + vocab_size = 23 + hidden_size = 8 + seq_len = 6 + cache_size = 12 + + teacher_heads = { + "0": torch.randn(vocab_size, hidden_size) / hidden_size**0.5, + "1": torch.randn(vocab_size, hidden_size) / hidden_size**0.5, + } + teacher_caches = { + "0": torch.randn(cache_size, hidden_size) / hidden_size**0.5, + "1": torch.randn(cache_size, hidden_size) / hidden_size**0.5, + } + teacher_files = make_teacher_files(tmp_path, teacher_heads, teacher_caches) + + runner = object.__new__(ModelRunner) + runner.train_config = {} + runner.lm_head_fp32 = True + runner.pp_enabled = False + runner._opd_head_manager = None + runner._opd_head_config = None + runner._opd_hidden_cache = None + runner._opd_hidden_config = None + + hidden_states = (torch.randn(1, seq_len, hidden_size, device=device) / hidden_size**0.5).requires_grad_(True) + student_weight = (torch.randn(vocab_size, hidden_size, device=device) / hidden_size**0.5).requires_grad_(True) + micro_batch = { + "labels": torch.tensor([[2, 3, 4, 5, 6, 7]], device=device), + "teacher_ids": torch.tensor([[0, 0, 0, 1, 1, 1]], device=device), + "teacher_cache_indices": torch.tensor([[0, 1, 2, 3, 4, 5]], device=device), + "teacher_weights": torch.tensor([[1.0, 0.5, 1.5, 1.0, 0.25, 2.0]], device=device), + } + + params = { + "teacher_heads": teacher_files.heads, + "teacher_hidden_caches": teacher_files.hidden_caches, + "num_chunks": 2, + } + optimizer = torch.optim.Adam([hidden_states, student_weight], lr=0.1) + loss_history: list[float] = [] + for _ in range(8): + optimizer.zero_grad(set_to_none=True) + result = runner._compute_opd_micro_batch_loss( + hidden_states=hidden_states, + student_weight=student_weight, + micro_batch=micro_batch, + params=params, + ) + assert result.loss.isfinite() + assert result.metrics["valid_tokens"] == seq_len + assert result.metrics["opd_num_teachers"] == 2 + loss_history.append(float(result.loss.detach().cpu())) + result.loss.backward() + assert hidden_states.grad is not None and hidden_states.grad.isfinite().all() + assert student_weight.grad is not None and student_weight.grad.isfinite().all() + optimizer.step() + + assert loss_history[-1] < loss_history[0], f"OPD loss did not decrease: {loss_history}" diff --git a/tests/e2e/test_opd_sglang_teacher.py b/tests/e2e/test_opd_sglang_teacher.py new file mode 100644 index 00000000..b9ff4638 --- /dev/null +++ b/tests/e2e/test_opd_sglang_teacher.py @@ -0,0 +1,245 @@ +"""End-to-end OPD coverage with a real xorl-sglang teacher.""" + +import json +import math +import os +import subprocess +import textwrap +from pathlib import Path + +import pytest +import torch + +from .e2e_utils import create_tiny_model_dir +from .server_utils import ( + ServerProcess, + _get_free_port, + _start_server_or_fail, + generate_server_config, + get_sglang_paths, + post_and_wait_for_future, + scalar, +) + + +SGLANG_PYTHON, SGLANG_SOURCE_DIR = get_sglang_paths() + + +pytestmark = [ + pytest.mark.e2e, + pytest.mark.gpu, + pytest.mark.server, + pytest.mark.slow, + pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"), + pytest.mark.skipif( + SGLANG_PYTHON is None or not SGLANG_PYTHON.exists(), + reason="Set XORL_SGLANG_PYTHON or XORL_SGLANG_INTERNAL_DIR", + ), + pytest.mark.skipif( + SGLANG_SOURCE_DIR is None or not SGLANG_SOURCE_DIR.exists(), + reason="Set XORL_SGLANG_SOURCE_DIR or XORL_SGLANG_INTERNAL_DIR", + ), +] + + +def _visible_teacher_gpu() -> str: + visible = os.environ.get("CUDA_VISIBLE_DEVICES") + if visible: + return visible.split(",")[-1].strip() + return str(max(torch.cuda.device_count() - 1, 0)) + + +def _run_sglang_teacher_cache( + model_dir: str, token_sequences: list[list[int]], output_dir: Path +) -> tuple[Path, list[list[int]]]: + assert SGLANG_PYTHON is not None + assert SGLANG_SOURCE_DIR is not None + output_dir.mkdir(parents=True, exist_ok=True) + tokens_path = output_dir / "teacher_tokens.json" + hidden_path = output_dir / "teacher_hidden.safetensors" + metadata_path = output_dir / "teacher_cache_metadata.json" + script_path = output_dir / "build_teacher_cache.py" + + tokens_path.write_text(json.dumps(token_sequences), encoding="utf-8") + script_path.write_text( + textwrap.dedent( + """ + import json + import sys + from pathlib import Path + + import torch + import sglang as sgl + from safetensors.torch import save_file + + + def main(): + model_dir = sys.argv[1] + tokens_path = Path(sys.argv[2]) + hidden_path = Path(sys.argv[3]) + metadata_path = Path(sys.argv[4]) + + token_sequences = json.loads(tokens_path.read_text(encoding="utf-8")) + hidden_size = json.loads((Path(model_dir) / "config.json").read_text(encoding="utf-8"))["hidden_size"] + engine = sgl.Engine( + model_path=model_dir, + tokenizer_path=model_dir, + enable_return_hidden_states=True, + disable_cuda_graph=True, + disable_piecewise_cuda_graph=True, + disable_radix_cache=True, + attention_backend="triton", + dtype="bfloat16", + mem_fraction_static=0.2, + max_total_tokens=128, + log_level="error", + ) + + hidden_chunks = [] + cache_indices_by_sample = [] + offset = 0 + try: + for input_ids in token_sequences: + output = engine.generate( + input_ids=input_ids, + sampling_params={"temperature": 0, "max_new_tokens": 1}, + return_hidden_states=True, + ) + hidden_state_chunks = output["meta_info"]["hidden_states"] + if not hidden_state_chunks: + raise RuntimeError("SGLang did not return teacher hidden states") + hidden = torch.tensor(hidden_state_chunks[0], dtype=torch.bfloat16) + if hidden.shape != (len(input_ids), hidden_size): + raise RuntimeError( + f"Unexpected hidden state shape {tuple(hidden.shape)} for input length {len(input_ids)}" + ) + hidden_chunks.append(hidden.cpu()) + cache_indices_by_sample.append(list(range(offset, offset + len(input_ids)))) + offset += len(input_ids) + finally: + engine.shutdown() + + save_file({"hidden_states": torch.cat(hidden_chunks, dim=0).contiguous()}, str(hidden_path)) + metadata_path.write_text( + json.dumps({"cache_indices_by_sample": cache_indices_by_sample}, indent=2), + encoding="utf-8", + ) + + + if __name__ == "__main__": + main() + """ + ), + encoding="utf-8", + ) + + env = os.environ.copy() + env["PYTHONPATH"] = f"{SGLANG_SOURCE_DIR}:{env.get('PYTHONPATH', '')}" + env["CUDA_VISIBLE_DEVICES"] = _visible_teacher_gpu() + + result = subprocess.run( + [str(SGLANG_PYTHON), str(script_path), model_dir, str(tokens_path), str(hidden_path), str(metadata_path)], + env=env, + text=True, + capture_output=True, + timeout=240, + ) + if result.returncode != 0: + raise AssertionError( + "xorl-sglang teacher cache generation failed\n" + f"--- stdout ---\n{result.stdout[-4000:]}\n" + f"--- stderr ---\n{result.stderr[-4000:]}" + ) + + metadata = json.loads(metadata_path.read_text(encoding="utf-8")) + return hidden_path, metadata["cache_indices_by_sample"] + + +def test_opd_server_forward_backward_with_xorl_sglang_teacher(tmp_path): + torch.manual_seed(2026) + + student_model_dir = create_tiny_model_dir(str(tmp_path / "student"), model_type="dense", save_weights=True) + teacher_model_dir = create_tiny_model_dir(str(tmp_path / "teacher"), model_type="dense", save_weights=True) + + token_sequences = [ + [11, 12, 13, 14], + [21, 22, 23], + [31, 32, 33, 34, 35], + ] + teacher_hidden_path, cache_indices = _run_sglang_teacher_cache( + teacher_model_dir, + token_sequences, + tmp_path / "teacher_artifacts", + ) + + output_dir = tmp_path / "xorl_server" + config_path = generate_server_config( + student_model_dir, + str(output_dir), + num_gpus=1, + enable_lora=False, + enable_gradient_checkpointing=False, + sample_packing_sequence_len=32, + extra_config={ + "enable_full_shard": False, + "worker_connection_timeout": 180.0, + }, + ) + server = ServerProcess(config_path=config_path, num_gpus=1, api_port=_get_free_port(), output_dir=str(output_dir)) + + try: + _start_server_or_fail(server, timeout=240.0) + + data = [] + for tokens, indices in zip(token_sequences, cache_indices): + data.append( + { + "model_input": {"input_ids": tokens}, + "loss_fn_inputs": { + "target_tokens": tokens, + "teacher_ids": [0] * len(tokens), + "teacher_weights": [1.0] * len(tokens), + "teacher_cache_indices": indices, + }, + } + ) + + forward_backward = post_and_wait_for_future( + server.base_url, + "/api/v1/forward_backward", + { + "model_id": "default", + "forward_backward_input": { + "data": data, + "loss_fn": "opd_loss", + "loss_fn_params": { + "teacher_heads": {"0": teacher_model_dir}, + "teacher_hidden_caches": {"0": str(teacher_hidden_path)}, + "opd_sort_by_teacher": True, + "num_chunks": 8, + }, + }, + }, + request_timeout=30.0, + future_timeout=240.0, + ) + + loss = scalar(forward_backward["loss_fn_outputs"][0]["loss"]) + assert math.isfinite(loss) + assert loss >= 0.0 + assert forward_backward["metrics"]["valid_tokens:sum"] == sum(len(tokens) for tokens in token_sequences) + + optim_step = post_and_wait_for_future( + server.base_url, + "/api/v1/optim_step", + { + "model_id": "default", + "adam_params": {"learning_rate": 1e-4, "beta1": 0.9, "beta2": 0.95, "eps": 1e-8}, + "gradient_clip": 1.0, + }, + request_timeout=30.0, + future_timeout=240.0, + ) + assert math.isfinite(scalar(optim_step["metrics"]["grad_norm"])) + finally: + server.stop() diff --git a/tests/models/test_deepseek_v3_checkpoint_handler.py b/tests/models/test_deepseek_v3_checkpoint_handler.py new file mode 100644 index 00000000..be9d3a1f --- /dev/null +++ b/tests/models/test_deepseek_v3_checkpoint_handler.py @@ -0,0 +1,272 @@ +import json +import math + +import pytest +import torch + +from xorl.models.transformers.deepseek_v3.checkpoint_handler import DeepseekV3CheckpointHandler +from xorl.models.transformers.deepseek_v3.configuration_deepseek_v3 import DeepseekV3Config +from xorl.models.transformers.deepseek_v3.modeling_deepseek_v3 import DeepseekV3ForCausalLM + + +pytestmark = [pytest.mark.cpu] + + +def _expert_weight(expert_idx: int, proj: str) -> torch.Tensor: + hidden_size = 2 + intermediate_size = 3 + value = float(expert_idx * 10 + {"gate": 1, "up": 2, "down": 3}[proj]) + if proj == "down": + return torch.full((hidden_size, intermediate_size), value) + return torch.full((intermediate_size, hidden_size), value) + + +def _pack_int4(values: torch.Tensor) -> torch.Tensor: + if values.dtype != torch.int8: + raise ValueError(f"Expected int8 values to pack, got {values.dtype}") + if values.ndim != 2: + raise ValueError(f"Expected rank-2 tensor to pack, got {tuple(values.shape)}") + + num_bits = 4 + pack_factor = 32 // num_bits + unsigned = (values + (1 << (num_bits - 1))).to(torch.uint8) + pad_cols = (-values.shape[1]) % pack_factor + if pad_cols: + unsigned = torch.nn.functional.pad(unsigned, (0, pad_cols)) + reshaped = unsigned.view(values.shape[0], -1, pack_factor).to(torch.int32) + bit_shifts = torch.arange(pack_factor, dtype=torch.int32) * num_bits + return (reshaped << bit_shifts).sum(dim=2, dtype=torch.int32) + + +def _packed_expert_weight(expert_idx: int, proj: str) -> dict[str, torch.Tensor]: + dense_weight = _expert_weight(expert_idx, proj) + quantized = torch.ones_like(dense_weight, dtype=torch.int8) + num_groups = max(1, math.ceil(dense_weight.shape[1] / 32)) + scales = torch.full((dense_weight.shape[0], num_groups), dense_weight.flatten()[0].item(), dtype=torch.float32) + return { + "weight_packed": _pack_int4(quantized), + "weight_scale": scales, + "weight_shape": torch.tensor(dense_weight.shape, dtype=torch.int64), + } + + +def _tiny_config() -> DeepseekV3Config: + config = DeepseekV3Config( + vocab_size=32, + hidden_size=16, + intermediate_size=32, + moe_intermediate_size=8, + num_hidden_layers=1, + num_attention_heads=2, + num_key_value_heads=2, + n_shared_experts=1, + n_routed_experts=4, + routed_scaling_factor=1.0, + kv_lora_rank=4, + q_lora_rank=8, + qk_nope_head_dim=4, + qk_rope_head_dim=4, + v_head_dim=8, + n_group=2, + topk_group=1, + num_experts_per_tok=2, + first_k_dense_replace=0, + ) + config._attn_implementation = "eager" + config._activation_native = True + return config + + +def test_checkpoint_handler_merges_language_model_experts_and_skips_multimodal_keys(): + handler = DeepseekV3CheckpointHandler(num_experts=4) + loaded = {} + + for expert_idx in range(4): + for proj in ("gate", "up", "down"): + loaded.update( + handler.on_load_weight( + f"language_model.model.layers.0.mlp.experts.{expert_idx}.{proj}_proj.weight", + _expert_weight(expert_idx, proj), + ) + ) + + loaded.update(handler.on_load_weight("language_model.model.layers.0.self_attn.o_proj.weight", torch.eye(2))) + assert handler.on_load_weight("vision_tower.encoder.weight", torch.ones(1)) == [] + assert handler.on_load_weight("mm_projector.weight", torch.ones(1)) == [] + + gate_up = dict(loaded)["model.layers.0.mlp.experts.gate_up_proj"] + down = dict(loaded)["model.layers.0.mlp.experts.down_proj"] + + assert gate_up.shape == (4, 2, 6) + assert down.shape == (4, 3, 2) + assert torch.all(gate_up[0, :, :3] == 1.0) + assert torch.all(gate_up[0, :, 3:] == 2.0) + assert torch.all(gate_up[3, :, :3] == 31.0) + assert torch.all(gate_up[3, :, 3:] == 32.0) + assert torch.all(down[1] == 13.0) + assert torch.equal(dict(loaded)["model.layers.0.self_attn.o_proj.weight"], torch.eye(2)) + + +def test_checkpoint_handler_splits_internal_fused_experts_on_save(): + handler = DeepseekV3CheckpointHandler(num_experts=2) + gate = torch.arange(2 * 2 * 3, dtype=torch.float32).reshape(2, 2, 3) + up = gate + 100.0 + gate_up = torch.cat([gate, up], dim=2) + down = torch.arange(2 * 3 * 2, dtype=torch.float32).reshape(2, 3, 2) + + split_gate_up = dict(handler.on_save_weight("model.layers.0.mlp.experts.gate_up_proj", gate_up)) + split_down = dict(handler.on_save_weight("model.layers.0.mlp.experts.down_proj", down)) + + assert torch.equal(split_gate_up["model.layers.0.mlp.experts.0.gate_proj.weight"], gate[0].transpose(0, 1)) + assert torch.equal(split_gate_up["model.layers.0.mlp.experts.1.up_proj.weight"], up[1].transpose(0, 1)) + assert torch.equal(split_down["model.layers.0.mlp.experts.0.down_proj.weight"], down[0].transpose(0, 1)) + assert torch.equal(split_down["model.layers.0.mlp.experts.1.down_proj.weight"], down[1].transpose(0, 1)) + + +def test_checkpoint_handler_keeps_internal_fused_expert_layout_on_load(): + handler = DeepseekV3CheckpointHandler(num_experts=2) + gate_up = torch.arange(2 * 2 * 6, dtype=torch.float32).reshape(2, 2, 6) + down = torch.arange(2 * 3 * 2, dtype=torch.float32).reshape(2, 3, 2) + + loaded_gate_up = dict(handler.on_load_weight("model.layers.0.mlp.experts.gate_up_proj", gate_up)) + loaded_down = dict(handler.on_load_weight("model.layers.0.mlp.experts.down_proj", down)) + + assert torch.equal(loaded_gate_up["model.layers.0.mlp.experts.gate_up_proj"], gate_up) + assert torch.equal(loaded_down["model.layers.0.mlp.experts.down_proj"], down) + + +def test_checkpoint_handler_ep_slices_to_local_experts(): + handler = DeepseekV3CheckpointHandler(num_experts=4, ep_rank=1, ep_size=2) + skip_key = handler.get_skip_key_fn() + loaded = {} + + for expert_idx in range(4): + for proj in ("gate", "up", "down"): + key = f"language_model.model.layers.0.mlp.experts.{expert_idx}.{proj}_proj.weight" + if skip_key is not None and skip_key(key): + loaded.update(handler.on_skip_weight(key)) + else: + loaded.update(handler.on_load_weight(key, _expert_weight(expert_idx, proj))) + + gate_up = dict(loaded)["model.layers.0.mlp.experts.gate_up_proj"] + down = dict(loaded)["model.layers.0.mlp.experts.down_proj"] + + assert gate_up.shape == (2, 2, 6) + assert down.shape == (2, 3, 2) + assert gate_up[:, 0, 0].tolist() == [21.0, 31.0] + assert down[:, 0, 0].tolist() == [23.0, 33.0] + + +def test_checkpoint_handler_loads_packed_expert_weights(): + handler = DeepseekV3CheckpointHandler(num_experts=4) + loaded = {} + + for expert_idx in range(4): + for proj in ("gate", "up", "down"): + for suffix, tensor in _packed_expert_weight(expert_idx, proj).items(): + loaded.update( + handler.on_load_weight( + f"language_model.model.layers.0.mlp.experts.{expert_idx}.{proj}_proj.{suffix}", + tensor, + ) + ) + + gate_up = dict(loaded)["model.layers.0.mlp.experts.gate_up_proj"] + down = dict(loaded)["model.layers.0.mlp.experts.down_proj"] + + assert gate_up.shape == (4, 2, 6) + assert down.shape == (4, 3, 2) + assert torch.all(gate_up[0, :, :3] == 1.0) + assert torch.all(gate_up[0, :, 3:] == 2.0) + assert torch.all(gate_up[3, :, :3] == 31.0) + assert torch.all(gate_up[3, :, 3:] == 32.0) + assert torch.all(down[1] == 13.0) + + +def test_checkpoint_handler_loads_packed_expert_weights_in_requested_dtype(): + handler = DeepseekV3CheckpointHandler(num_experts=4, device=torch.device("cpu"), dtype=torch.bfloat16) + loaded = {} + + for expert_idx in range(4): + for proj in ("gate", "up", "down"): + for suffix, tensor in _packed_expert_weight(expert_idx, proj).items(): + loaded.update( + handler.on_load_weight( + f"language_model.model.layers.0.mlp.experts.{expert_idx}.{proj}_proj.{suffix}", + tensor, + ) + ) + + gate_up = dict(loaded)["model.layers.0.mlp.experts.gate_up_proj"] + down = dict(loaded)["model.layers.0.mlp.experts.down_proj"] + + assert handler._expert_buffer is not None + assert handler._expert_buffer._device == torch.device("cpu") + assert gate_up.dtype == torch.bfloat16 + assert down.dtype == torch.bfloat16 + assert torch.all(gate_up[0, :, :3] == torch.tensor(1.0, dtype=torch.bfloat16)) + assert torch.all(down[1] == torch.tensor(13.0, dtype=torch.bfloat16)) + + +def test_checkpoint_handler_ep_slices_packed_experts_to_local_experts(): + handler = DeepseekV3CheckpointHandler(num_experts=4, ep_rank=1, ep_size=2) + skip_key = handler.get_skip_key_fn() + loaded = {} + + for expert_idx in range(4): + for proj in ("gate", "up", "down"): + for suffix, tensor in _packed_expert_weight(expert_idx, proj).items(): + key = f"language_model.model.layers.0.mlp.experts.{expert_idx}.{proj}_proj.{suffix}" + if skip_key is not None and skip_key(key): + loaded.update(handler.on_skip_weight(key)) + else: + loaded.update(handler.on_load_weight(key, tensor)) + + gate_up = dict(loaded)["model.layers.0.mlp.experts.gate_up_proj"] + down = dict(loaded)["model.layers.0.mlp.experts.down_proj"] + + assert gate_up.shape == (2, 2, 6) + assert down.shape == (2, 3, 2) + assert gate_up[:, 0, 0].tolist() == [21.0, 31.0] + assert down[:, 0, 0].tolist() == [23.0, 33.0] + + +def test_model_checkpoint_handler_accepts_official_packed_expert_layout(): + model = DeepseekV3ForCausalLM(_tiny_config()) + handler = model.get_checkpoint_handler( + checkpoint_keys={"language_model.model.layers.0.mlp.experts.0.gate_proj.weight_packed"}, + ) + assert isinstance(handler, DeepseekV3CheckpointHandler) + + +def test_model_checkpoint_handler_reads_packed_quant_config_from_text_config(tmp_path): + (tmp_path / "config.json").write_text( + json.dumps( + { + "model_type": "kimi_k25", + "text_config": { + "quantization_config": { + "quant_method": "compressed-tensors", + "format": "pack-quantized", + "config_groups": { + "group_0": { + "weights": { + "group_size": 64, + "num_bits": 8, + } + } + }, + } + }, + } + ) + ) + + model = DeepseekV3ForCausalLM(_tiny_config()) + handler = model.get_checkpoint_handler( + checkpoint_keys={"language_model.model.layers.0.mlp.experts.0.gate_proj.weight_packed"}, + weights_path=str(tmp_path), + ) + + assert handler._packed_expert_group_size == 64 + assert handler._packed_expert_num_bits == 8 diff --git a/tests/models/test_deepseek_v3_model.py b/tests/models/test_deepseek_v3_model.py new file mode 100644 index 00000000..1182e3c2 --- /dev/null +++ b/tests/models/test_deepseek_v3_model.py @@ -0,0 +1,146 @@ +import pytest +import torch + +from xorl.lora.modules import LoraLinear +from xorl.lora.utils import inject_lora_into_model +from xorl.models.layers.moe import MoEExpertsLoRA +from xorl.models.layers.moe.routing_replay import RoutingReplay, set_replay_stage +from xorl.models.transformers.deepseek_v3.configuration_deepseek_v3 import DeepseekV3Config +from xorl.models.transformers.deepseek_v3.modeling_deepseek_v3 import DeepseekV3ForCausalLM +from xorl.models.transformers.deepseek_v3.support import freeze_deepseek_v3_router_parameters + + +pytestmark = [pytest.mark.cpu] + + +def _tiny_config() -> DeepseekV3Config: + config = DeepseekV3Config( + vocab_size=32, + hidden_size=16, + intermediate_size=32, + moe_intermediate_size=8, + num_hidden_layers=2, + num_attention_heads=2, + num_key_value_heads=2, + n_shared_experts=1, + n_routed_experts=4, + routed_scaling_factor=1.0, + kv_lora_rank=4, + q_lora_rank=8, + qk_nope_head_dim=4, + qk_rope_head_dim=4, + v_head_dim=8, + n_group=2, + topk_group=1, + num_experts_per_tok=2, + first_k_dense_replace=0, + max_position_embeddings=64, + rope_theta=10000.0, + attention_dropout=0.0, + topk_method="noaux_tc", + scoring_func="sigmoid", + ) + config._attn_implementation = "eager" + config._moe_implementation = "eager" + config._activation_native = True + return config + + +def test_deepseek_v3_tiny_forward_backward_and_freeze_router(): + model = DeepseekV3ForCausalLM(_tiny_config()) + model.train() + + input_ids = torch.randint(0, model.config.vocab_size, (2, 5)) + outputs = model( + input_ids=input_ids, + attention_mask=torch.ones_like(input_ids), + output_router_logits=True, + ) + + assert tuple(outputs.last_hidden_state.shape) == (2, 5, model.config.hidden_size) + assert len(outputs.router_logits) == model.config.num_hidden_layers + + loss = outputs.last_hidden_state.float().sum() + loss.backward() + + assert model.model.layers[0].self_attn.q_a_proj.weight.grad is not None + assert model.model.layers[0].mlp.shared_experts.down_proj.weight.grad is not None + + frozen = freeze_deepseek_v3_router_parameters(model) + + assert frozen == model.config.num_hidden_layers + for name, param in model.named_parameters(): + if ".gate.weight" in name: + assert param.requires_grad is False + + +def test_deepseek_v3_default_lora_targets_cover_mla_and_moe(): + model = DeepseekV3ForCausalLM(_tiny_config()) + + inject_lora_into_model(model, r=4, lora_alpha=8, target_modules=None) + + attn = model.model.layers[0].self_attn + mlp = model.model.layers[0].mlp + + assert isinstance(attn.q_a_proj, LoraLinear) + assert isinstance(attn.q_b_proj, LoraLinear) + assert isinstance(attn.kv_a_proj_with_mqa, LoraLinear) + assert isinstance(attn.kv_b_proj, LoraLinear) + assert isinstance(attn.o_proj, LoraLinear) + assert isinstance(mlp.shared_experts.gate_proj, LoraLinear) + assert isinstance(mlp.shared_experts.up_proj, LoraLinear) + assert isinstance(mlp.shared_experts.down_proj, LoraLinear) + assert isinstance(mlp.experts, MoEExpertsLoRA) + + +def test_deepseek_v3_forward_emits_router_logits_when_aux_loss_is_enabled_by_config(): + config = _tiny_config() + config.output_router_logits = False + config.router_aux_loss_coef = 0.001 + + model = DeepseekV3ForCausalLM(config) + outputs = model( + input_ids=torch.randint(0, model.config.vocab_size, (2, 5)), + attention_mask=torch.ones(2, 5, dtype=torch.long), + ) + + assert outputs.router_logits is not None + assert len(outputs.router_logits) == model.config.num_hidden_layers + + +def test_deepseek_v3_router_logits_skip_dense_layers_for_aux_loss(): + config = _tiny_config() + config.first_k_dense_replace = 1 + config.output_router_logits = False + config.router_aux_loss_coef = 0.001 + + model = DeepseekV3ForCausalLM(config) + outputs = model( + input_ids=torch.randint(0, model.config.vocab_size, (2, 5)), + attention_mask=torch.ones(2, 5, dtype=torch.long), + ) + + assert outputs.router_logits is not None + assert len(outputs.router_logits) == model.config.num_hidden_layers - config.first_k_dense_replace + assert all(router_logits is not None for router_logits in outputs.router_logits) + + +def test_deepseek_v3_routing_replay_records_weights(): + model = DeepseekV3ForCausalLM(_tiny_config()) + block = model.model.layers[0].mlp + replay = RoutingReplay() + block._routing_replay = replay + hidden_states = torch.randn(2, 3, model.config.hidden_size) + + try: + set_replay_stage("record") + block(hidden_states) + finally: + set_replay_stage(None) + + assert len(replay.top_indices_list) == 1 + assert len(replay.top_weights_list) == 1 + assert tuple(replay.top_weights_list[0].shape) == ( + hidden_states.shape[0] * hidden_states.shape[1], + model.config.num_experts_per_tok, + ) diff --git a/tests/models/test_deepseek_v3_registry.py b/tests/models/test_deepseek_v3_registry.py new file mode 100644 index 00000000..f0fd584d --- /dev/null +++ b/tests/models/test_deepseek_v3_registry.py @@ -0,0 +1,125 @@ +import json +from types import SimpleNamespace + +import pytest + +from xorl.models.auto import _load_local_xorl_config +from xorl.models.registry import get_registry +from xorl.models.transformers.deepseek_v3.configuration_deepseek_v3 import DeepseekV3Config +from xorl.models.transformers.deepseek_v3.modeling_deepseek_v3 import DeepseekV3ForCausalLM + + +pytestmark = [pytest.mark.cpu] + + +def _make_kimi_text_config(): + return dict( + model_type="kimi_k2", + architectures=["DeepseekV3ForCausalLM"], + vocab_size=163840, + hidden_size=128, + intermediate_size=256, + moe_intermediate_size=32, + num_hidden_layers=4, + num_attention_heads=4, + num_key_value_heads=4, + kv_lora_rank=16, + q_lora_rank=32, + qk_nope_head_dim=16, + qk_rope_head_dim=8, + v_head_dim=16, + n_routed_experts=8, + n_shared_experts=2, + n_group=4, + topk_group=2, + num_experts_per_tok=2, + first_k_dense_replace=1, + routed_scaling_factor=1.25, + norm_topk_prob=True, + hidden_act="silu", + max_position_embeddings=4096, + initializer_range=0.02, + rms_norm_eps=1e-6, + use_cache=True, + attention_bias=False, + attention_dropout=0.0, + output_router_logits=True, + router_aux_loss_coef=0.001, + topk_method="noaux_tc", + scoring_func="sigmoid", + rope_scaling={"rope_type": "default", "rope_theta": 1000000.0}, + ) + + +def test_deepseek_v3_registered(): + registry = get_registry() + assert "DeepseekV3ForCausalLM" in registry.supported_models + assert registry.get_model_cls_from_model_arch("DeepseekV3ForCausalLM") is DeepseekV3ForCausalLM + + +def test_deepseek_v3_config_from_kimi_wrapper_hf_config(): + hf_config = SimpleNamespace( + model_type="kimi_k25", + text_config=SimpleNamespace(**_make_kimi_text_config()), + vision_config=SimpleNamespace(hidden_size=1024, num_hidden_layers=24), + tie_word_embeddings=False, + ) + + config = DeepseekV3Config.from_hf_config(hf_config) + + assert config.model_type == "deepseek_v3" + assert config.q_lora_rank == 32 + assert config.kv_lora_rank == 16 + assert config.qk_nope_head_dim == 16 + assert config.qk_rope_head_dim == 8 + assert config.v_head_dim == 16 + assert config.n_routed_experts == 8 + assert config.n_shared_experts == 2 + assert config.first_k_dense_replace == 1 + assert config.topk_method == "noaux_tc" + assert config.scoring_func == "sigmoid" + assert config.rope_scaling["rope_type"] == "default" + assert config.rope_theta == 1000000.0 + + +def test_deepseek_v3_config_maps_official_kimi_aux_loss_defaults(): + kimi_text_config = _make_kimi_text_config() + kimi_text_config.pop("router_aux_loss_coef") + kimi_text_config.pop("output_router_logits") + kimi_text_config["aux_loss_alpha"] = 0.001 + + hf_config = SimpleNamespace( + model_type="kimi_k25", + text_config=SimpleNamespace(**kimi_text_config), + tie_word_embeddings=False, + ) + + config = DeepseekV3Config.from_hf_config(hf_config) + + assert config.router_aux_loss_coef == pytest.approx(0.001) + assert config.output_router_logits is True + + +def test_local_auto_config_unwraps_kimi_wrapper_text_config(tmp_path): + config_dir = tmp_path / "kimi-k25" + config_dir.mkdir() + (config_dir / "config.json").write_text( + json.dumps( + { + "model_type": "kimi_k25", + "text_config": _make_kimi_text_config(), + "vision_config": { + "hidden_size": 1024, + "num_hidden_layers": 24, + }, + "tie_word_embeddings": False, + } + ) + ) + + config = _load_local_xorl_config(str(config_dir), {}) + + assert isinstance(config, DeepseekV3Config) + assert config.model_type == "deepseek_v3" + assert config.n_routed_experts == 8 + assert config.n_shared_experts == 2 diff --git a/tests/models/test_gradient_checkpointing.py b/tests/models/test_gradient_checkpointing.py new file mode 100644 index 00000000..c9debfe4 --- /dev/null +++ b/tests/models/test_gradient_checkpointing.py @@ -0,0 +1,92 @@ +"""Tests for `GradientCheckpointingLayer` + `gradient_checkpointing_enable` contract.""" + +from unittest.mock import MagicMock + +import pytest +import torch + +from xorl.models.base import XorlPreTrainedModel +from xorl.models.module_utils import ( + DEFAULT_GRADIENT_CHECKPOINTING_METHOD, + GradientCheckpointingLayer, + MoEGradientCheckpointingLayer, +) + + +pytestmark = [pytest.mark.cpu] + + +class _IdentityCheckpointLayer(GradientCheckpointingLayer): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x + + +class _StubMoECheckpointLayer(MoEGradientCheckpointingLayer): + """MoE stub for attribute-level assertions — forward is not exercised.""" + + +@pytest.fixture +def model() -> XorlPreTrainedModel: + m = XorlPreTrainedModel(config=None) + m.layer = _IdentityCheckpointLayer() + return m + + +@pytest.mark.parametrize( + "layer_cls", + [GradientCheckpointingLayer, MoEGradientCheckpointingLayer], +) +def test_class_default_method_is_the_recompute_default(layer_cls): + assert layer_cls._gradient_checkpointing_method == DEFAULT_GRADIENT_CHECKPOINTING_METHOD + + +@pytest.mark.parametrize( + "layer_cls", + [_IdentityCheckpointLayer, _StubMoECheckpointLayer], +) +def test_gradient_checkpointing_enable_default(layer_cls): + model = XorlPreTrainedModel(config=None) + model.layer = layer_cls() + model.gradient_checkpointing_enable() + + assert model.layer.gradient_checkpointing is True + assert model.layer._gradient_checkpointing_method == DEFAULT_GRADIENT_CHECKPOINTING_METHOD + + +@pytest.mark.parametrize( + "method", + ["recompute_full_layer", "recompute_before_dispatch", "no_recompute"], +) +def test_enable_propagates_method_kwarg_to_every_checkpointed_layer(method): + model = XorlPreTrainedModel(config=None) + model.layer = _IdentityCheckpointLayer() + model.moe_layer = _StubMoECheckpointLayer() + + model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"gradient_checkpointing_method": method}) + + assert model.layer._gradient_checkpointing_method == method + assert model.moe_layer._gradient_checkpointing_method == method + + +@pytest.mark.parametrize( + "training, flag_enabled, expect_checkpoint", + [ + (True, True, True), + (True, False, False), + (False, True, False), + (False, False, False), + ], +) +def test_outer_gate_fires_iff_training_and_flag(model, training, flag_enabled, expect_checkpoint): + model.gradient_checkpointing_enable() + model.train(training) + model.layer.gradient_checkpointing = flag_enabled + + spy = MagicMock(side_effect=lambda fn, *a, **kw: fn(*a, **kw)) + model.layer._gradient_checkpointing_func = spy + + x = torch.zeros(2, 3) + out = model.layer(x) + + assert torch.equal(out, x) + assert spy.called is expect_checkpoint diff --git a/tests/models/test_kimi_tokenizer.py b/tests/models/test_kimi_tokenizer.py new file mode 100644 index 00000000..cea1d48f --- /dev/null +++ b/tests/models/test_kimi_tokenizer.py @@ -0,0 +1,82 @@ +import json + +import pytest +from tiktoken.load import dump_tiktoken_bpe + +from xorl.models import auto as auto_module +from xorl.models.auto import build_processor, build_tokenizer +from xorl.models.transformers.deepseek_v3.tokenization_kimi import TikTokenTokenizer + + +pytestmark = [pytest.mark.cpu] + + +def _write_tiktoken_fixture(tmp_path): + tokenizer_dir = tmp_path / "kimi-tokenizer" + tokenizer_dir.mkdir() + dump_tiktoken_bpe({bytes([i]): i for i in range(256)}, str(tokenizer_dir / "tiktoken.model")) + (tokenizer_dir / "tokenizer_config.json").write_text( + json.dumps( + { + "tokenizer_class": "TikTokenTokenizer", + "auto_map": {"AutoTokenizer": ["tokenization_kimi.TikTokenTokenizer", None]}, + "bos_token": "[BOS]", + "eos_token": "[EOS]", + "unk_token": "[UNK]", + "pad_token": "[PAD]", + "additional_special_tokens": ["<|im_end|>"], + "added_tokens_decoder": { + "256": {"content": "[BOS]", "special": True}, + "257": {"content": "[EOS]", "special": True}, + "258": {"content": "[UNK]", "special": True}, + "259": {"content": "[PAD]", "special": True}, + "260": {"content": "<|im_end|>", "special": True}, + }, + } + ) + ) + return tokenizer_dir + + +def test_build_tokenizer_loads_local_kimi_tiktoken_without_remote_code(tmp_path): + tokenizer_dir = _write_tiktoken_fixture(tmp_path) + + tokenizer = build_tokenizer(str(tokenizer_dir)) + + assert isinstance(tokenizer, TikTokenTokenizer) + assert tokenizer.bos_token_id == 256 + assert tokenizer.eos_token_id == 257 + assert tokenizer.pad_token_id == 259 + assert tokenizer.decode(tokenizer.encode("hello")) == "hello" + + +def test_build_tokenizer_fallback_does_not_enable_remote_code(monkeypatch, tmp_path): + calls = {} + + def fake_from_pretrained(path, **kwargs): + calls["path"] = path + calls["kwargs"] = kwargs + return object() + + monkeypatch.setattr(auto_module.AutoTokenizer, "from_pretrained", fake_from_pretrained) + + build_tokenizer(str(tmp_path / "not-kimi")) + + assert "trust_remote_code" not in calls["kwargs"] + assert calls["kwargs"]["padding_side"] == "right" + + +def test_build_processor_does_not_enable_remote_code(monkeypatch): + calls = {} + + def fake_from_pretrained(path, **kwargs): + calls["path"] = path + calls["kwargs"] = kwargs + return object() + + monkeypatch.setattr(auto_module.AutoProcessor, "from_pretrained", fake_from_pretrained) + + build_processor("processor-path") + + assert "trust_remote_code" not in calls["kwargs"] + assert calls["kwargs"]["padding_side"] == "right" diff --git a/tests/models/test_lora_merge_fp32_cast_once.py b/tests/models/test_lora_merge_fp32_cast_once.py new file mode 100644 index 00000000..49968d8c --- /dev/null +++ b/tests/models/test_lora_merge_fp32_cast_once.py @@ -0,0 +1,140 @@ +"""Tests for the fp32 cast-once merge variant on both LoraLinear and MoEExpertsLoRA. + +Invariants: + - With zero LoRA (B=0), merge must be bit-exact: W_merged == W (no change). + - After merge, ``merged_weight`` equals the fp32 reference + ``(W.to(fp32) + B@A*s).to(W.dtype)`` bit-for-bit (by construction). + - Merged weight is >= as faithful as the naive ``W + Δ.to(W.dtype)`` variant — + i.e., the fp32-sum-then-cast distance to the true fp32 merged value is ≤ + the naive-merge distance. +""" + +import pytest +import torch + + +pytestmark = [pytest.mark.gpu] + + +def _naive_merge(weight, delta): + """Old behavior for comparison: round Δ per-element, then add.""" + return weight + delta.to(weight.dtype) + + +def _fp32_merge(weight, delta): + """New behavior: add in fp32, cast once.""" + return (weight.to(torch.float32) + delta).to(weight.dtype) + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32]) +def test_lora_linear_merge_zero_b_is_bitexact(dtype): + """merge_weights() on a fresh LoraLinear (B=0) must leave weight untouched.""" + from xorl.lora import LoraLinear + + torch.manual_seed(0) + layer = LoraLinear(128, 64, r=8, lora_alpha=16, device="cuda", dtype=dtype) + layer.weight.data.normal_(std=0.05) + before = layer.weight.detach().clone() + layer.merge_weights() + assert torch.equal(layer.weight, before), "zero-LoRA merge must be bit-exact" + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +def test_lora_linear_merge_matches_fp32_reference(dtype): + """The merged weight must equal ``(W.to(fp32) + B@A*s).to(W.dtype)`` exactly.""" + from xorl.lora import LoraLinear + + torch.manual_seed(0) + layer = LoraLinear(128, 64, r=8, lora_alpha=16, device="cuda", dtype=dtype) + layer.weight.data.normal_(std=0.05) + # non-zero LoRA + layer.lora_B.data.normal_(std=0.02) + w_before = layer.weight.detach().clone() + delta_fp32 = (layer.lora_B @ layer.lora_A) * layer.scaling + expected = _fp32_merge(w_before, delta_fp32) + + layer.merge_weights() + assert torch.equal(layer.weight, expected), "merge_weights must match fp32-cast-once reference" + + +def test_lora_linear_merge_strictly_ge_naive_precision(): + """For any B, fp32-cast-once ≤ naive in distance to true fp32 merged value.""" + from xorl.lora import LoraLinear + + torch.manual_seed(42) + layer = LoraLinear(128, 64, r=8, lora_alpha=16, device="cuda", dtype=torch.bfloat16) + layer.weight.data.normal_(std=0.05) + layer.lora_B.data.normal_(std=0.02) + + w = layer.weight.detach().clone() + delta_fp32 = (layer.lora_B @ layer.lora_A) * layer.scaling + true_fp32 = w.to(torch.float32) + delta_fp32 # reference in fp32 + naive = _naive_merge(w, delta_fp32).to(torch.float32) # rounds Δ then adds + fp32_cast_once = _fp32_merge(w, delta_fp32).to(torch.float32) # fp32 sum, cast once + + naive_err = (naive - true_fp32).abs().max().item() + fp32_err = (fp32_cast_once - true_fp32).abs().max().item() + assert fp32_err <= naive_err + 1e-9, ( + f"fp32-cast-once should be ≤ naive precision: fp32={fp32_err:.3e} naive={naive_err:.3e}" + ) + + +def _tiny_moe_experts_with_lora(dtype): + from xorl.lora import MoEExpertsLoRA, MoELoRAConfig + + cfg = MoELoRAConfig(r=8, lora_alpha=16, target_modules=["gate_proj", "up_proj", "down_proj"]) + e = ( + MoEExpertsLoRA( + num_experts=4, + hidden_dim=32, + intermediate_size=24, + hidden_act="silu", + moe_implementation="eager", + lora_config=cfg, + ) + .to(dtype) + .cuda() + ) + e.gate_up_proj.data.normal_(std=0.05) + e.down_proj.data.normal_(std=0.05) + return e + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +def test_moe_merge_zero_b_is_bitexact(dtype): + e = _tiny_moe_experts_with_lora(dtype) + gu_before = e.gate_up_proj.detach().clone() + dn_before = e.down_proj.detach().clone() + e.merge_weights() + assert torch.equal(e.gate_up_proj, gu_before), "zero-LoRA MoE merge must not change gate_up_proj" + assert torch.equal(e.down_proj, dn_before), "zero-LoRA MoE merge must not change down_proj" + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +def test_moe_merge_matches_fp32_reference(dtype): + e = _tiny_moe_experts_with_lora(dtype) + # perturb all lora_B + torch.manual_seed(7) + for name in ("gate_proj", "up_proj", "down_proj"): + getattr(e, f"{name}_lora_B").data.normal_(std=0.02) + + # expected = fp32 sum then cast, computed from SNAPSHOTS of current state + gu_before = e.gate_up_proj.detach().clone() + dn_before = e.down_proj.detach().clone() + expected_updates = {} + for proj in ("gate_proj", "up_proj", "down_proj"): + delta = e._compute_proj_delta(proj) # fp32, [E, in, out] + expected_updates[proj] = delta + + # gate/up land in the fused gate_up_proj via views, each shape (E, H, I) + I = e.intermediate_size + gate_expected = gu_before.clone() + gate_expected[..., :I] = _fp32_merge(gu_before[..., :I], expected_updates["gate_proj"]) + gate_expected[..., I:] = _fp32_merge(gu_before[..., I:], expected_updates["up_proj"]) + + down_expected = _fp32_merge(dn_before, expected_updates["down_proj"]) + + e.merge_weights() + + assert torch.equal(e.gate_up_proj, gate_expected), "MoE gate+up merge must match fp32 reference" + assert torch.equal(e.down_proj, down_expected), "MoE down merge must match fp32 reference" diff --git a/tests/models/test_lora_moe_attention_targets.py b/tests/models/test_lora_moe_attention_targets.py new file mode 100644 index 00000000..0bb54e82 --- /dev/null +++ b/tests/models/test_lora_moe_attention_targets.py @@ -0,0 +1,70 @@ +"""Regression test for inject_lora_into_model_with_moe attention partition. + +Locks in the fix for the case where the function silently dropped DeepSeek-V3 / +Kimi MLA attention projections (q_a_proj, q_b_proj, kv_a_proj_with_mqa, +kv_b_proj) because the attention vs. expert split was hardcoded to a +Llama/Qwen-shaped allowlist (q_proj/k_proj/v_proj/o_proj/lm_head). +""" + +import pytest +import torch.nn as nn + +from xorl.lora import LoraLinear, inject_lora_into_model_with_moe + + +pytestmark = [pytest.mark.cpu] + + +class _StubConfig: + def __init__(self, model_type: str): + self.model_type = model_type + self.num_experts = 0 + + +class _MLALikeBlock(nn.Module): + """Single attention block with DeepSeek-V3 / Kimi MLA projection names.""" + + def __init__(self, hidden_size: int = 32, q_lora_rank: int = 16, kv_lora_rank: int = 16): + super().__init__() + self.q_a_proj = nn.Linear(hidden_size, q_lora_rank, bias=False) + self.q_b_proj = nn.Linear(q_lora_rank, hidden_size, bias=False) + self.kv_a_proj_with_mqa = nn.Linear(hidden_size, kv_lora_rank, bias=False) + self.kv_b_proj = nn.Linear(kv_lora_rank, hidden_size, bias=False) + self.o_proj = nn.Linear(hidden_size, hidden_size, bias=False) + + +class _DeepSeekLikeModel(nn.Module): + def __init__(self): + super().__init__() + self.config = _StubConfig("deepseek_v3") + self.self_attn = _MLALikeBlock() + + +def test_default_targets_cover_all_mla_projections_for_deepseek_v3(): + model = _DeepSeekLikeModel() + + inject_lora_into_model_with_moe(model, r=4, lora_alpha=8, target_modules=None) + + for proj in ("q_a_proj", "q_b_proj", "kv_a_proj_with_mqa", "kv_b_proj", "o_proj"): + replaced = getattr(model.self_attn, proj) + assert isinstance(replaced, LoraLinear), ( + f"{proj} was not LoRA-replaced; the attention/expert split is dropping MLA projections again." + ) + + +def test_explicit_mla_targets_are_not_filtered_out(): + model = _DeepSeekLikeModel() + + inject_lora_into_model_with_moe( + model, + r=4, + lora_alpha=8, + target_modules=["q_a_proj", "q_b_proj", "kv_a_proj_with_mqa", "kv_b_proj"], + ) + + for proj in ("q_a_proj", "q_b_proj", "kv_a_proj_with_mqa", "kv_b_proj"): + replaced = getattr(model.self_attn, proj) + assert isinstance(replaced, LoraLinear), ( + f"{proj} was filtered out of attention_modules even though the caller passed it explicitly." + ) + assert not isinstance(model.self_attn.o_proj, LoraLinear) diff --git a/tests/models/test_model_state.py b/tests/models/test_model_state.py new file mode 100644 index 00000000..4ccd5934 --- /dev/null +++ b/tests/models/test_model_state.py @@ -0,0 +1,59 @@ +from types import SimpleNamespace + +import pytest +import torch + +from xorl.checkpoint import checkpointer + + +pytestmark = [pytest.mark.cpu] + + +class _TinyModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.linear = torch.nn.Linear(4, 2, bias=False) + self.register_buffer("persistent_buf", torch.ones(3)) + self.register_buffer("scratch_buf", torch.zeros(2), persistent=False) + + +def test_reference_state_dict_bypasses_dcp_state_dict_and_skips_nonpersistent_buffers(monkeypatch): + monkeypatch.setattr(checkpointer, "get_parallel_state", lambda: SimpleNamespace(dp_mode="none")) + monkeypatch.setattr( + checkpointer, + "get_model_state_dict", + lambda *args, **kwargs: (_ for _ in ()).throw(AssertionError("unexpected DCP state_dict call")), + ) + + model_state = checkpointer.ModelState(_TinyModel()) + state_dict = model_state.reference_state_dict() + + assert "linear.weight" in state_dict + assert "persistent_buf" in state_dict + assert "scratch_buf" not in state_dict + + +def test_distributed_checkpointer_load_skips_missing_optimizer_state(tmp_path, monkeypatch): + captured = {} + + class _FakeReader: + def __init__(self, path): + self.path = path + + def read_metadata(self): + return SimpleNamespace(state_dict_metadata={"model.linear.weight": object()}) + + def fake_dcp_load(state_dict, storage_reader, process_group=None): + captured["state_keys"] = set(state_dict) + captured["storage_reader"] = storage_reader + captured["process_group"] = process_group + + monkeypatch.setattr(checkpointer, "FileSystemReader", _FakeReader) + monkeypatch.setattr(checkpointer.dcp, "load", fake_dcp_load) + + state = {"model": _TinyModel(), "optimizer": object()} + result = checkpointer.DistributedCheckpointer.load(str(tmp_path), state) + + assert result is state + assert captured["state_keys"] == {"model"} + assert isinstance(captured["storage_reader"], _FakeReader) diff --git a/tests/models/test_module_utils_broadcast.py b/tests/models/test_module_utils_broadcast.py index bb4681c5..e9cd2815 100644 --- a/tests/models/test_module_utils_broadcast.py +++ b/tests/models/test_module_utils_broadcast.py @@ -1,7 +1,15 @@ +import contextlib +import socket from types import SimpleNamespace import pytest import torch +import torch.distributed as dist +import torch.multiprocessing as mp +from torch.distributed._tensor import Replicate +from torch.distributed._tensor import Shard as DTShard +from torch.distributed.device_mesh import DeviceMesh +from torch.distributed.tensor import DTensor from xorl.models import module_utils @@ -20,6 +28,273 @@ def to_empty(self, device): self.device = device +class _FakeDeviceMesh: + ndim = 1 + + def __init__(self, size: int, local_rank: int): + self._size = size + self._local_rank = local_rank + + def size(self): + return self._size + + def get_local_rank(self): + return self._local_rank + + +class _FakeDTensor: + def __init__(self, local_tensor: torch.Tensor, mesh_size: int, local_rank: int, placements): + self._local_tensor = local_tensor + self.device_mesh = _FakeDeviceMesh(mesh_size, local_rank) + self.placements = placements + + +def _find_free_port() -> int: + with contextlib.closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as sock: + sock.bind(("127.0.0.1", 0)) + return int(sock.getsockname()[1]) + + +def _cpu_dtensor_materialize_worker(rank: int, world_size: int, port: int) -> None: + dist.init_process_group("gloo", init_method=f"tcp://127.0.0.1:{port}", rank=rank, world_size=world_size) + try: + module_utils._cpu_save_device_mesh_cache.clear() + mesh = DeviceMesh( + "cpu", + mesh=torch.arange(world_size).view(2, 2), + mesh_dim_names=("ep", "fsdp"), + backend_override=(("gloo", None), ("gloo", None)), + ) + full_tensor = torch.arange(16, dtype=torch.float32).view(4, 4) + row = rank // 2 + col = rank % 2 + local_tensor = full_tensor[row * 2 : (row + 1) * 2, col * 2 : (col + 1) * 2].clone() + dtensor = DTensor.from_local( + local_tensor, + device_mesh=mesh, + placements=[DTShard(0), DTShard(1)], + shape=full_tensor.shape, + stride=full_tensor.stride(), + ) + materialized = module_utils._materialize_tensor_for_save(dtensor) + assert materialized.device.type == "cpu" + assert torch.equal(materialized, full_tensor) + finally: + dist.destroy_process_group() + + +def _cpu_dtensor_materialize_to_rank_worker(rank: int, world_size: int, port: int, dst_rank: int) -> None: + dist.init_process_group("gloo", init_method=f"tcp://127.0.0.1:{port}", rank=rank, world_size=world_size) + try: + module_utils._cpu_save_device_mesh_cache.clear() + mesh = DeviceMesh( + "cpu", + mesh=torch.arange(world_size), + mesh_dim_names=("fsdp",), + backend_override=(("gloo", None),), + ) + full_tensor = torch.arange(18, dtype=torch.float32).view(9, 2) + chunk_size = 3 + start = rank * chunk_size + stop = min(start + chunk_size, full_tensor.shape[0]) + local_tensor = full_tensor[start:stop].clone() + dtensor = DTensor.from_local( + local_tensor, + device_mesh=mesh, + placements=[DTShard(0)], + shape=full_tensor.shape, + stride=full_tensor.stride(), + ) + materialized = module_utils._materialize_tensor_for_save(dtensor, dst_rank=dst_rank) + if rank == dst_rank: + assert materialized is not None + assert materialized.device.type == "cpu" + assert torch.equal(materialized, full_tensor) + else: + assert materialized is None + finally: + dist.destroy_process_group() + + +def _cpu_dtensor_materialize_2d_to_rank_worker(rank: int, world_size: int, port: int, dst_rank: int) -> None: + dist.init_process_group("gloo", init_method=f"tcp://127.0.0.1:{port}", rank=rank, world_size=world_size) + try: + module_utils._cpu_save_device_mesh_cache.clear() + mesh = DeviceMesh( + "cpu", + mesh=torch.arange(world_size).view(2, 2), + mesh_dim_names=("ep", "fsdp"), + backend_override=(("gloo", None), ("gloo", None)), + ) + full_tensor = torch.arange(16, dtype=torch.float32).view(4, 4) + row = rank // 2 + col = rank % 2 + local_tensor = full_tensor[row * 2 : (row + 1) * 2, col * 2 : (col + 1) * 2].clone() + dtensor = DTensor.from_local( + local_tensor, + device_mesh=mesh, + placements=[DTShard(0), DTShard(1)], + shape=full_tensor.shape, + stride=full_tensor.stride(), + ) + materialized = module_utils._materialize_tensor_for_save(dtensor, dst_rank=dst_rank) + if rank == dst_rank: + assert materialized is not None + assert materialized.device.type == "cpu" + assert torch.equal(materialized, full_tensor) + else: + assert materialized is None + finally: + dist.destroy_process_group() + + +def test_copy_into_existing_dtensor_shard_for_replicated_tensor(): + dtensor = _FakeDTensor(torch.zeros(4, dtype=torch.float32), mesh_size=4, local_rank=2, placements=(Replicate(),)) + full_tensor = torch.arange(4, dtype=torch.float32) + + copied = module_utils._copy_into_existing_dtensor_shard(dtensor, full_tensor) + + assert copied is True + assert torch.equal(dtensor._local_tensor, full_tensor) + + +def test_materialize_tensor_for_save_uses_cpu_mesh_for_dtensors(): + port = _find_free_port() + mp.start_processes( + _cpu_dtensor_materialize_worker, + args=(4, port), + nprocs=4, + join=True, + start_method="fork", + ) + + +def test_materialize_tensor_for_save_gathers_1d_dtensor_to_writer_rank(): + port = _find_free_port() + mp.start_processes( + _cpu_dtensor_materialize_to_rank_worker, + args=(4, port, 2), + nprocs=4, + join=True, + start_method="fork", + ) + + +def test_materialize_tensor_for_save_gathers_2d_dtensor_to_writer_rank(): + port = _find_free_port() + mp.start_processes( + _cpu_dtensor_materialize_2d_to_rank_worker, + args=(4, port, 3), + nprocs=4, + join=True, + start_method="fork", + ) + + +def test_copy_into_existing_dtensor_shard_for_sharded_tensor(): + dtensor = _FakeDTensor( + torch.zeros(2, 3, dtype=torch.float32), + mesh_size=4, + local_rank=1, + placements=(DTShard(0),), + ) + full_tensor = torch.arange(24, dtype=torch.float32).view(8, 3) + + copied = module_utils._copy_into_existing_dtensor_shard(dtensor, full_tensor) + + assert copied is True + assert torch.equal(dtensor._local_tensor, full_tensor[2:4]) + + +def test_copy_into_existing_dtensor_shard_trims_padded_tail_shards(): + dtensor = _FakeDTensor( + torch.zeros(0, 3, dtype=torch.float32), + mesh_size=8, + local_rank=6, + placements=(DTShard(0),), + ) + full_tensor = torch.arange(15, dtype=torch.float32).view(5, 3) + + copied = module_utils._copy_into_existing_dtensor_shard(dtensor, full_tensor) + + assert copied is True + assert tuple(dtensor._local_tensor.shape) == (0, 3) + + +def test_copy_into_existing_dtensor_shard_rejects_shape_mismatched_replicates(): + dtensor = _FakeDTensor(torch.zeros(1, 3, dtype=torch.float32), mesh_size=4, local_rank=0, placements=(Replicate(),)) + full_tensor = torch.arange(15, dtype=torch.float32).view(5, 3) + + copied = module_utils._copy_into_existing_dtensor_shard(dtensor, full_tensor) + + assert copied is False + assert torch.equal(dtensor._local_tensor, torch.zeros(1, 3, dtype=torch.float32)) + + +def test_broadcast_object_list_serializes_over_tensor_broadcast_for_nccl_groups(monkeypatch): + fake_group = object() + state = {"rank": 3} + stored = [] + + def fake_broadcast(tensor, src=0, group=None): + assert group is fake_group + if state["rank"] == src: + stored.append(tensor.detach().cpu().clone()) + else: + tensor.copy_(stored.pop(0).to(tensor.device)) + + fake_dist = SimpleNamespace( + get_rank=lambda: state["rank"], + broadcast=fake_broadcast, + broadcast_object_list=lambda *args, **kwargs: (_ for _ in ()).throw( + AssertionError("unexpected object broadcast") + ), + ) + + monkeypatch.setattr(module_utils, "dist", fake_dist) + monkeypatch.setattr(module_utils, "_get_object_broadcast_device", lambda group: torch.device("cpu")) + + source_payload = [[("payload", torch.Size([2, 3]), torch.float32)]] + module_utils._broadcast_object_list(source_payload, src=3, group=fake_group) + + state["rank"] = 1 + received_payload = [None] + module_utils._broadcast_object_list(received_payload, src=3, group=fake_group) + + assert received_payload == source_payload + + +def test_get_object_broadcast_device_uses_default_nccl_group(monkeypatch): + fake_dist = SimpleNamespace( + is_available=lambda: True, + is_initialized=lambda: True, + get_backend=lambda group=None: "nccl", + ) + + monkeypatch.setattr(module_utils, "dist", fake_dist) + monkeypatch.setattr(module_utils, "get_device_type", lambda: "cuda") + monkeypatch.setattr(module_utils, "get_device_id", lambda: 3) + + assert module_utils._get_object_broadcast_device(None) == torch.device("cuda:3") + + +def test_broadcast_object_list_weight_load_uses_weight_load_group(monkeypatch): + fake_group = object() + calls = [] + + monkeypatch.setattr(module_utils, "_get_weight_load_group", lambda: fake_group) + monkeypatch.setattr( + module_utils, + "_broadcast_object_list", + lambda obj_list, src=0, group=None: calls.append((obj_list, src, group)), + ) + + payload = [["checkpoint-paths"]] + module_utils._broadcast_object_list_weight_load(payload, src=7) + + assert calls == [(payload, 7, fake_group)] + + def test_rank0_broadcast_path_calls_load_state_dict_on_nonzero_ranks(monkeypatch): calls = [] @@ -36,7 +311,7 @@ def fake_broadcast_object_list(obj, src=0, group=None, device=None): 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_object_broadcast_device", lambda group: None) monkeypatch.setattr( module_utils, "get_parallel_state", @@ -114,7 +389,7 @@ def fail_prefetch(*args, **kwargs): 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_object_broadcast_device", lambda group: None) monkeypatch.setattr( module_utils, "get_parallel_state", @@ -163,11 +438,543 @@ def fake_try_load_state_dict_local(weights_path, **kwargs): 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, "_get_object_broadcast_device", lambda group: None) 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"] + + +@pytest.mark.parametrize( + ("key", "expected"), + [ + ("model.layers.43.mlp.experts.7.gate_proj.weight", True), + ("model.layers.43.mlp.experts.7.down_proj.weight_scale_inv", True), + ("model.layers.43.mlp.experts.gate_up_proj", True), + ("model.layers.43.mlp.experts.gate_up_proj.weight", True), + ("model.layers.43.mlp.experts.down_proj.weight", True), + ("language_model.model.layers.43.mlp.experts.down_proj.weight", True), + ("layers.12.ffn.experts.7.w1.weight", True), + ("layers.12.ffn.experts.7.w2.scale", True), + ("model.layers.12.ffn.experts.7.w3.weight", True), + ("model.layers.43.mlp.shared_expert.down_proj.weight", False), + ("layers.12.ffn.shared_experts.w1.weight", False), + ("model.layers.43.mlp.gate_up_proj.weight", False), + ("model.layers.43.self_attn.q_proj.weight", False), + ], +) +def test_checkpoint_expert_key_classifies_supported_expert_formats(key, expected): + assert module_utils._is_checkpoint_expert_key(key) is expected + + +def test_try_load_state_dict_local_directory_skips_broadcast(monkeypatch): + local_resolution_calls = [] + + fake_dist = SimpleNamespace( + is_initialized=lambda: True, + get_rank=lambda: 7, + get_world_size=lambda: 64, + broadcast_object_list=lambda *args, **kwargs: (_ for _ in ()).throw(AssertionError("should not broadcast")), + ) + + def fake_try_load_state_dict_local(weights_path, **kwargs): + local_resolution_calls.append(weights_path) + return [module_utils.StateDictIterator("local-shard.safetensors")] + + monkeypatch.setattr(module_utils, "dist", fake_dist) + monkeypatch.setattr(module_utils.os.path, "isdir", lambda path: path == "dummy-local-dir") + monkeypatch.setattr(module_utils, "_try_load_state_dict_local", fake_try_load_state_dict_local) + + iterators = module_utils._try_load_state_dict("dummy-local-dir") + + assert local_resolution_calls == ["dummy-local-dir"] + assert [it.filepath for it in iterators] == ["local-shard.safetensors"] + + +def test_checkpoint_expert_filter_handles_wrapped_language_model_keys(): + assert module_utils._is_checkpoint_expert_key("model.language_model.layers.43.mlp.experts.gate_up_proj") + assert module_utils._is_checkpoint_expert_key("model.language_model.layers.43.mlp.experts.down_proj") + assert not module_utils._is_checkpoint_expert_key( + "model.language_model.layers.43.mlp.shared_expert.down_proj.weight" + ) + + +def test_grouped_load_weights_uses_filtered_prefetch_on_group_leader(monkeypatch): + batch_meta_calls = [] + dispatched = [] + transfer_calls = [] + handler_kwargs = [] + handler_calls = {"dense_loaded": [], "expert_loaded": []} + fake_group = object() + fake_dense_group = object() + + class _DenseHandler: + def get_skip_key_fn(self): + return None + + def on_load_weight(self, key, tensor): + handler_calls["dense_loaded"].append(key) + return [(key, tensor)] + + def on_load_complete(self): + return [] + + class _ExpertHandler: + def get_skip_key_fn(self): + return None + + def on_load_weight(self, key, tensor): + handler_calls["expert_loaded"].append(key) + return [("model.layers.0.mlp.experts.gate_proj", torch.arange(8, dtype=torch.float32).view(8, 1, 1))] + + def on_load_complete(self): + return [] + + class _GroupedModel: + def named_buffers(self): + return [] + + def named_parameters(self): + return [ + ("keep.weight", None), + ("model.layers.0.mlp.experts.gate_proj", None), + ] + + def named_modules(self): + return [] + + def to_empty(self, device): + self.device = device + + def get_checkpoint_handler(self, **kwargs): + handler_kwargs.append(kwargs) + return _DenseHandler() if kwargs["ep_size"] == 1 else _ExpertHandler() + + def fake_broadcast_object_list(obj, src=0, group=None, device=None): + batch_meta_calls.append((src, group, obj[0])) + if obj[0] is None: + obj[0] = [] + + fake_dist = SimpleNamespace( + is_available=lambda: True, + is_initialized=lambda: True, + get_world_size=lambda: 128, + get_process_group_ranks=lambda group: ( + [0, 1, 2, 3, 4, 5, 6, 7] if group is fake_dense_group else [0, 16, 32, 48, 64, 80, 96, 112] + ), + broadcast=lambda tensor, src=0, group=None: transfer_calls.append( + ("broadcast", src, group, tuple(tensor.shape)) + ), + scatter=lambda tensor, scatter_list=None, src=0, group=None: transfer_calls.append( + ("scatter", src, group, tuple(tensor.shape), None if scatter_list is None else tuple(scatter_list[0].shape)) + ), + broadcast_object_list=fake_broadcast_object_list, + ) + + prefetch_calls = [] + + def fake_prefetch_filtered(state_dict_iterators, skip_key_fn, prefetch_count): + assert state_dict_iterators == ["shard-0"] + if skip_key_fn("model.layers.0.mlp.experts.0.gate_proj.weight"): + prefetch_calls.append(("dense", prefetch_count)) + assert prefetch_count == 1 + assert not skip_key_fn("keep.weight") + yield ({"keep.weight": torch.tensor([2.0])}, []) + return + + prefetch_calls.append(("expert", prefetch_count)) + assert prefetch_count == 1 + assert skip_key_fn("keep.weight") + yield ({"model.layers.0.mlp.experts.0.gate_proj.weight": torch.tensor([[3.0]])}, []) + + 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_object_broadcast_device", lambda group: None) + monkeypatch.setattr(module_utils, "_get_grouped_weight_load_group", lambda _ps: fake_group) + monkeypatch.setattr(module_utils, "_get_grouped_dense_weight_load_group", lambda: fake_dense_group) + monkeypatch.setattr( + module_utils, + "get_parallel_state", + lambda: SimpleNamespace(global_rank=0, pp_enabled=False, ep_enabled=True, ep_rank=0, ep_size=16), + ) + 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: {"keep.weight", "model.layers.0.mlp.experts.0.gate_proj.weight"}, + ) + monkeypatch.setattr( + module_utils, + "_get_expert_scatter_target_shape", + lambda model, parameter_name, tensor, parallel_plan, parallel_state: ( + (1, 1, 1) if parameter_name == "model.layers.0.mlp.experts.gate_proj" else None + ), + ) + 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.grouped_load_weights(_GroupedModel(), "dummy-weights", init_device="cpu") + + assert handler_kwargs == [ + { + "checkpoint_keys": {"keep.weight", "model.layers.0.mlp.experts.0.gate_proj.weight"}, + "ep_rank": 0, + "ep_size": 1, + "is_broadcast": False, + "weights_path": "dummy-weights", + "device": None, + "dtype": None, + }, + { + "checkpoint_keys": {"keep.weight", "model.layers.0.mlp.experts.0.gate_proj.weight"}, + "ep_rank": 0, + "ep_size": 16, + "is_broadcast": False, + "weights_path": "dummy-weights", + "device": None, + "dtype": None, + }, + ] + assert prefetch_calls == [("dense", 1), ("expert", 1)] + assert handler_calls["dense_loaded"] == ["keep.weight"] + assert handler_calls["expert_loaded"] == ["model.layers.0.mlp.experts.0.gate_proj.weight"] + assert dispatched == ["keep.weight", "model.layers.0.mlp.experts.gate_proj"] + assert transfer_calls == [ + ("broadcast", 0, fake_dense_group, (1,)), + ("scatter", 0, fake_group, (1, 1, 1), (1, 1, 1)), + ] + assert batch_meta_calls[0] == ( + 0, + fake_dense_group, + [ + ("keep.weight", torch.Size([1]), torch.float32, "broadcast"), + ], + ) + assert batch_meta_calls[1] == ( + 0, + fake_group, + [ + ("model.layers.0.mlp.experts.gate_proj", torch.Size([1, 1, 1]), torch.float32, "expert_scatter"), + ], + ) + + +def test_grouped_load_weights_routes_hf_fused_experts_through_expert_queue(monkeypatch): + dense_loaded = [] + expert_loaded = [] + dispatched = [] + prefetch_calls = [] + fake_group = object() + fake_dense_group = object() + raw_expert_key = "model.language_model.layers.0.mlp.experts.gate_up_proj.weight" + raw_skipped_key = "mtp.layers.0.mlp.experts.0.gate_proj.weight" + expert_key = "model.layers.0.mlp.experts.gate_up_proj.weight" + expert_param = "model.layers.0.mlp.experts.gate_up_proj" + + class _Plan: + def is_expert_parameter(self, parameter_name): + return parameter_name == expert_param + + class _Handler: + def __init__(self, loaded): + self.loaded = loaded + + def get_skip_key_fn(self): + return None + + def on_load_weight(self, key, tensor): + self.loaded.append(key) + if key == expert_key: + return [(expert_param, tensor)] + return [(key, tensor)] + + def on_load_complete(self): + return [] + + class _GroupedModel: + _checkpoint_conversion_mapping = {r"^model\.language_model\.": "model."} + _checkpoint_skip_key_patterns = [r"^mtp\."] + + def named_buffers(self): + return [] + + def named_parameters(self): + return [("keep.weight", None), (expert_param, None)] + + def named_modules(self): + return [] + + def to_empty(self, device): + self.device = device + + def get_parallel_plan(self): + return _Plan() + + def get_checkpoint_handler(self, **kwargs): + return _Handler(dense_loaded if kwargs["ep_size"] == 1 else expert_loaded) + + fake_dist = SimpleNamespace( + is_available=lambda: True, + is_initialized=lambda: True, + get_world_size=lambda group=None: 1, + get_process_group_ranks=lambda group: [0], + broadcast=lambda *args, **kwargs: None, + scatter=lambda *args, **kwargs: None, + broadcast_object_list=lambda *args, **kwargs: None, + ) + + def fake_prefetch_filtered(state_dict_iterators, skip_key_fn, prefetch_count): + assert state_dict_iterators == ["shard-0"] + if skip_key_fn(raw_expert_key): + prefetch_calls.append("dense") + assert skip_key_fn(raw_skipped_key) + assert not skip_key_fn("keep.weight") + yield ({"keep.weight": torch.tensor([2.0])}, []) + else: + prefetch_calls.append("expert") + assert skip_key_fn(raw_skipped_key) + assert skip_key_fn("keep.weight") + yield ({raw_expert_key: torch.tensor([3.0])}, []) + + monkeypatch.setattr(module_utils, "dist", fake_dist) + monkeypatch.setattr(module_utils, "tqdm", lambda iterable, **kwargs: iterable) + monkeypatch.setattr(module_utils, "_get_object_broadcast_device", lambda group: None) + monkeypatch.setattr(module_utils, "_get_grouped_weight_load_group", lambda _ps: fake_group) + monkeypatch.setattr(module_utils, "_get_grouped_dense_weight_load_group", lambda: fake_dense_group) + monkeypatch.setattr( + module_utils, + "get_parallel_state", + lambda: SimpleNamespace(global_rank=0, pp_enabled=False, ep_enabled=True, ep_rank=0, ep_size=16), + ) + 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: {"keep.weight", raw_expert_key, raw_skipped_key}, + ) + 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, "_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.grouped_load_weights(_GroupedModel(), "dummy-weights", init_device="cpu") + + assert prefetch_calls == ["dense", "expert"] + assert dense_loaded == ["keep.weight"] + assert expert_loaded == [expert_key] + assert dispatched == ["keep.weight", expert_param] + + +def test_grouped_load_weights_routes_ffn_expert_source_format_through_expert_queue(monkeypatch): + dense_loaded = [] + expert_loaded = [] + dispatched = [] + prefetch_calls = [] + fake_group = object() + fake_dense_group = object() + raw_expert_weight = "layers.0.ffn.experts.0.w1.weight" + raw_expert_scale = "layers.0.ffn.experts.0.w1.scale" + expert_param = "model.layers.0.mlp.experts.gate_up_proj" + + class _Plan: + def is_expert_parameter(self, parameter_name): + return parameter_name == expert_param + + class _Handler: + def __init__(self, loaded): + self.loaded = loaded + + def get_skip_key_fn(self): + return None + + def on_load_weight(self, key, tensor): + self.loaded.append(key) + if key == raw_expert_weight: + return [(expert_param, tensor)] + return [(key, tensor)] if key == "keep.weight" else [] + + def on_load_complete(self): + return [] + + class _GroupedModel: + def named_buffers(self): + return [] + + def named_parameters(self): + return [("keep.weight", None), (expert_param, None)] + + def named_modules(self): + return [] + + def to_empty(self, device): + self.device = device + + def get_parallel_plan(self): + return _Plan() + + def get_checkpoint_handler(self, **kwargs): + return _Handler(dense_loaded if kwargs["ep_size"] == 1 else expert_loaded) + + fake_dist = SimpleNamespace( + is_available=lambda: True, + is_initialized=lambda: True, + get_world_size=lambda group=None: 1, + get_process_group_ranks=lambda group: [0], + broadcast=lambda *args, **kwargs: None, + scatter=lambda *args, **kwargs: None, + broadcast_object_list=lambda *args, **kwargs: None, + ) + + def fake_prefetch_filtered(state_dict_iterators, skip_key_fn, prefetch_count): + assert state_dict_iterators == ["shard-0"] + if skip_key_fn(raw_expert_weight): + prefetch_calls.append("dense") + assert skip_key_fn(raw_expert_scale) + assert not skip_key_fn("keep.weight") + yield ({"keep.weight": torch.tensor([2.0])}, []) + else: + prefetch_calls.append("expert") + assert skip_key_fn("keep.weight") + assert not skip_key_fn(raw_expert_scale) + yield ({raw_expert_weight: torch.tensor([3.0]), raw_expert_scale: torch.tensor([1.0])}, []) + + monkeypatch.setattr(module_utils, "dist", fake_dist) + monkeypatch.setattr(module_utils, "tqdm", lambda iterable, **kwargs: iterable) + monkeypatch.setattr(module_utils, "_get_object_broadcast_device", lambda group: None) + monkeypatch.setattr(module_utils, "_get_grouped_weight_load_group", lambda _ps: fake_group) + monkeypatch.setattr(module_utils, "_get_grouped_dense_weight_load_group", lambda: fake_dense_group) + monkeypatch.setattr( + module_utils, + "get_parallel_state", + lambda: SimpleNamespace(global_rank=0, pp_enabled=False, ep_enabled=True, ep_rank=0, ep_size=16), + ) + 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: {"keep.weight", raw_expert_weight, raw_expert_scale}, + ) + 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, "_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.grouped_load_weights(_GroupedModel(), "dummy-weights", init_device="cpu") + + assert prefetch_calls == ["dense", "expert"] + assert dense_loaded == ["keep.weight"] + assert expert_loaded == [raw_expert_weight, raw_expert_scale] + assert dispatched == ["keep.weight", expert_param] + + +def test_grouped_load_weights_treats_missing_dense_group_as_local(monkeypatch): + dispatched = [] + fake_group = object() + + class _Handler: + def get_skip_key_fn(self): + return None + + def on_load_weight(self, key, tensor): + return [(key, tensor)] + + def on_load_complete(self): + return [] + + class _GroupedModel: + def named_buffers(self): + return [] + + def named_parameters(self): + return [("keep.weight", None)] + + def named_modules(self): + return [] + + def to_empty(self, device): + self.device = device + + def get_checkpoint_handler(self, **kwargs): + return _Handler() + + def fail_collective(*args, **kwargs): + raise AssertionError("local-only grouped dense loading should not enter a collective") + + fake_dist = SimpleNamespace( + is_available=lambda: True, + is_initialized=lambda: True, + get_world_size=lambda group=None: 4 if group is None else 1, + get_process_group_ranks=lambda group: [0], + broadcast=fail_collective, + scatter=fail_collective, + broadcast_object_list=fail_collective, + ) + + def fake_prefetch_filtered(state_dict_iterators, skip_key_fn, prefetch_count): + if skip_key_fn("keep.weight"): + yield ({}, []) + else: + yield ({"keep.weight": torch.tensor([2.0])}, []) + + monkeypatch.setattr(module_utils, "dist", fake_dist) + monkeypatch.setattr(module_utils, "tqdm", lambda iterable, **kwargs: iterable) + monkeypatch.setattr(module_utils, "_get_grouped_weight_load_group", lambda _ps: fake_group) + monkeypatch.setattr(module_utils, "_get_grouped_dense_weight_load_group", lambda: None) + monkeypatch.setattr( + module_utils, + "get_parallel_state", + lambda: SimpleNamespace(global_rank=0, pp_enabled=False, ep_enabled=True, ep_rank=0, ep_size=4), + ) + 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: {"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, "_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.grouped_load_weights(_GroupedModel(), "dummy-weights", init_device="cpu") + + assert dispatched == ["keep.weight"] + + +def test_grouped_load_weights_falls_back_without_ep_group(monkeypatch): + called = [] + + fake_dist = SimpleNamespace( + is_available=lambda: True, + is_initialized=lambda: True, + ) + + monkeypatch.setattr(module_utils, "dist", fake_dist) + monkeypatch.setattr(module_utils, "_get_grouped_weight_load_group", lambda _ps: None) + monkeypatch.setattr( + module_utils, + "get_parallel_state", + lambda: SimpleNamespace(global_rank=0, pp_enabled=False, ep_enabled=False, ep_fsdp_device_mesh=None), + ) + monkeypatch.setattr( + module_utils, + "rank0_load_and_broadcast_weights", + lambda *args, **kwargs: called.append((args, kwargs)), + ) + + model = _DummyModel() + module_utils.grouped_load_weights(model, "dummy-weights", init_device="cpu") + + assert len(called) == 1 + assert not hasattr(model, "device") diff --git a/tests/models/test_moe_experts_lora.py b/tests/models/test_moe_experts_lora.py index b88a9c76..8a006e89 100644 --- a/tests/models/test_moe_experts_lora.py +++ b/tests/models/test_moe_experts_lora.py @@ -1,14 +1,13 @@ """Tests for MoE experts with LoRA across all backends (eager, triton, native, quack).""" -import pytest - -from xorl.lora import LoraLinear, inject_lora_into_model +from dataclasses import make_dataclass +from unittest.mock import MagicMock, patch - -pytestmark = [pytest.mark.cpu, pytest.mark.gpu] +import pytest 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 ( @@ -17,6 +16,9 @@ ) +pytestmark = [pytest.mark.cpu, pytest.mark.gpu] + + class MockConfig: """Mock config for testing.""" @@ -162,6 +164,25 @@ def test_init_frozen_trainable_shapes(self, backend): assert f"num_experts={config.num_experts}" in repr_str assert f"r={lora_config.r}" in repr_str + def test_runtime_rank_lora_views_are_contiguous(self): + """Partial-rank views passed to group GEMM backends must be contiguous.""" + config = MockConfig() + experts = MoEExpertsLoRA( + num_experts=config.num_experts, + hidden_dim=config.hidden_size, + intermediate_size=config.moe_intermediate_size, + lora_config=MoELoRAConfig(r=8, lora_alpha=16), + ) + + experts.set_runtime_lora_config(lora_rank=3, lora_alpha=12) + + for proj_name in ("gate_proj", "up_proj", "down_proj"): + lora_A, lora_B = experts._active_lora_views(proj_name) + assert lora_A.shape[-1] == 3 + assert lora_B.shape[1] == 3 + assert lora_A.is_contiguous() + assert lora_B.is_contiguous() + # --------------------------------------------------------------------------- # 3. Eager LoRA forward/backward (CPU) @@ -707,9 +728,6 @@ def _run_ep_forward(self, experts, score_attr=None, scores=None, compute_output= 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) @@ -785,9 +803,6 @@ def test_no_scores_leaves_output_unchanged(self): 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) diff --git a/tests/models/test_moe_train_router_dispatch.py b/tests/models/test_moe_train_router_dispatch.py index b62377fc..a587a66c 100644 --- a/tests/models/test_moe_train_router_dispatch.py +++ b/tests/models/test_moe_train_router_dispatch.py @@ -6,6 +6,7 @@ from xorl.arguments import ModelArguments from xorl.models.layers.moe.moe_block import MoEBlock +from xorl.models.layers.moe.router import TopKRouter from xorl.server.server_arguments import ServerArguments @@ -77,3 +78,48 @@ def test_from_config_defaults_train_router_false(): moe = MoEBlock.from_config(config, moe_implementation="eager") assert moe.train_router is False + + +def test_balanced_synthetic_routing_env(monkeypatch): + monkeypatch.setenv("XORL_MOE_SYNTHETIC_ROUTING", "balanced") + router = TopKRouter(num_experts=4, top_k=2) + + logits = torch.randn(8, 4) + routing_weights, selected_experts = router(logits, torch.bfloat16) + + expected_experts = torch.tensor( + [ + [0, 1], + [2, 3], + [0, 1], + [2, 3], + [0, 1], + [2, 3], + [0, 1], + [2, 3], + ] + ) + assert torch.equal(selected_experts, expected_experts) + assert routing_weights.dtype == torch.bfloat16 + torch.testing.assert_close(routing_weights.float(), torch.full((8, 2), 0.5)) + + counts = torch.bincount(selected_experts.flatten(), minlength=4) + assert torch.equal(counts, torch.full((4,), 4, dtype=counts.dtype)) + + +def test_balanced_synthetic_routing_replay_regather_uses_uniform_weights(monkeypatch): + monkeypatch.setenv("XORL_MOE_SYNTHETIC_ROUTING", "balanced") + moe = MoEBlock( + hidden_size=16, + num_experts=4, + top_k=2, + intermediate_size=32, + moe_implementation="eager", + ) + + router_logits = torch.randn(3, 4) + cached_experts = torch.tensor([[3, 2], [1, 0], [2, 1]]) + selected_experts, routing_weights = moe._regather_routing(router_logits, cached_experts, torch.float32) + + assert torch.equal(selected_experts, cached_experts) + torch.testing.assert_close(routing_weights, torch.full((3, 2), 0.5)) diff --git a/tests/models/test_olmo2_support.py b/tests/models/test_olmo2_support.py new file mode 100644 index 00000000..9791dd6f --- /dev/null +++ b/tests/models/test_olmo2_support.py @@ -0,0 +1,180 @@ +import pytest +import torch +from torch.distributed.tensor.parallel import ColwiseParallel, RowwiseParallel +from transformers.models.olmo2.configuration_olmo2 import Olmo2Config as HFOlmo2Config +from transformers.models.olmo2.modeling_olmo2 import Olmo2ForCausalLM as HFOlmo2ForCausalLM + +from xorl.models.auto import build_foundation_model +from xorl.models.transformers.olmo2.configuration_olmo2 import Olmo2Config as XOlmo2Config +from xorl.models.transformers.olmo2.modeling_olmo2 import Olmo2ForCausalLM +from xorl.models.transformers.olmo2.parallelize import MODEL_TP_PLAN, TP_PLAN +from xorl.models.transformers.olmo2.tp_styles import LocalAxisRMSNormShard + + +pytestmark = [pytest.mark.cpu] + + +_COMMON_KWARGS = dict( + architectures=["Olmo2ForCausalLM"], + 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, + rope_theta=500000.0, + # OLMo-2 is post-norm: RMSNorm is applied to attn/MLP output. With the + # default 0.02 init the activations are tiny and RMSNorm amplifies fp32 + # rounding so much that HF and xorl drift apart numerically. A larger + # init keeps the comparison in the regime where rms ≫ eps. + initializer_range=0.5, + attention_dropout=0.0, + tie_word_embeddings=False, + use_cache=False, +) + + +def _make_hf_olmo2_config(): + return HFOlmo2Config(**_COMMON_KWARGS) + + +def _make_xorl_olmo2_config(): + config = XOlmo2Config(**_COMMON_KWARGS) + config._attn_implementation = "eager" + config._activation_native = True + return config + + +def test_build_foundation_model_accepts_hf_olmo2_config_object(): + hf_config = _make_hf_olmo2_config() + + model = build_foundation_model(hf_config, init_device="meta", attn_implementation="eager") + + assert isinstance(model, Olmo2ForCausalLM) + assert model.config.model_type == "olmo2" + layer = model.model.layers[0] + # OLMo-2 uses post-norm: no input_layernorm, has post-attn and post-feedforward norms. + assert not hasattr(layer, "input_layernorm") + assert hasattr(layer, "post_attention_layernorm") + assert hasattr(layer, "post_feedforward_layernorm") + # Full-axis QK norms (head_dim * num_heads), not per-head. + head_dim = hf_config.hidden_size // hf_config.num_attention_heads + assert layer.self_attn.q_norm.weight.shape == (hf_config.num_attention_heads * head_dim,) + assert layer.self_attn.k_norm.weight.shape == (hf_config.num_key_value_heads * head_dim,) + assert layer.self_attn.qkv_proj.bias is None + assert layer.self_attn.o_proj.bias is None + + +def test_olmo2_tp_plan_uses_local_axis_qk_norm(): + # OLMo-2's full-axis q_norm/k_norm doesn't compose with SequenceParallel + # or stock ColwiseParallel — under colwise q/k_proj the input arrives + # hidden-sharded and a full-hidden weight can't be applied directly. + # LocalAxisRMSNormShard shards the 1-D weight on dim 0 so each rank's + # slice matches its local q/k slice. This is the actual root cause of + # what was originally reported (the issue thought it was post-norm; the trace + # was at q_norm in _project_qkv). + assert isinstance(TP_PLAN["layers.*.self_attn.q_norm"], LocalAxisRMSNormShard) + assert isinstance(TP_PLAN["layers.*.self_attn.k_norm"], LocalAxisRMSNormShard) + + # Post-norms see a Replicate input after the rowwise all-reduce in + # o_proj/down_proj — they should NOT be in the plan (no TP wrapping). + assert "layers.*.post_attention_layernorm" not in TP_PLAN + assert "layers.*.post_feedforward_layernorm" not in TP_PLAN + assert "norm" not in TP_PLAN + + # Standard colwise/rowwise everywhere else. + assert isinstance(TP_PLAN["layers.*.self_attn.q_proj"], ColwiseParallel) + assert isinstance(TP_PLAN["layers.*.self_attn.o_proj"], RowwiseParallel) + assert isinstance(TP_PLAN["layers.*.mlp.gate_proj"], ColwiseParallel) + assert isinstance(TP_PLAN["layers.*.mlp.down_proj"], RowwiseParallel) + + # lm_head: vanilla colwise (Replicate input from the model's final norm, + # vocab-parallel output for vocab_parallel_cross_entropy). + assert MODEL_TP_PLAN["lm_head"] == "colwise" or isinstance(MODEL_TP_PLAN["lm_head"], ColwiseParallel) + + +def test_olmo2_unfuse_for_tp_matches_hf_parameter_layout(): + model = Olmo2ForCausalLM(_make_xorl_olmo2_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 None + assert layer.self_attn.k_proj.bias is None + assert layer.self_attn.v_proj.bias is 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_olmo2_checkpoint_handler_exports_hf_compatible_attention_keys(): + model = Olmo2ForCausalLM(_make_xorl_olmo2_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.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.self_attn.q_norm.weight" in transformed + assert "model.layers.0.self_attn.k_norm.weight" in transformed + assert "model.layers.0.post_attention_layernorm.weight" in transformed + assert "model.layers.0.post_feedforward_layernorm.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.qkv_proj.weight" not in transformed + assert "model.layers.0.mlp.gate_up_proj.weight" not in transformed + + +def test_olmo2_checkpoint_handler_loads_hf_weights_into_fused_model(): + hf_config = _make_hf_olmo2_config() + hf_config._attn_implementation = "eager" + xorl_config = _make_xorl_olmo2_config() + + hf_model = HFOlmo2ForCausalLM(hf_config) + xorl_model = Olmo2ForCausalLM(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.mlp.gate_up_proj.weight" in transformed + assert "model.layers.0.self_attn.q_norm.weight" in transformed + assert "model.layers.0.self_attn.k_norm.weight" in transformed + + load_result = xorl_model.load_state_dict(transformed, strict=False) + assert not load_result.missing_keys + assert not load_result.unexpected_keys + + # Avoid pad_token_id (1) so the embedded sequence has real activations + # at every position; otherwise tiny attn outputs amplify in post-norm. + input_ids = torch.tensor([[2, 3, 4, 5]]) + 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=1e-4, rtol=1e-4) + torch.testing.assert_close(xorl_logits, hf_logits, atol=2e-4, rtol=5e-4) diff --git a/tests/models/test_qwen3_5_apply_rotary.py b/tests/models/test_qwen3_5_apply_rotary.py new file mode 100644 index 00000000..1182016b --- /dev/null +++ b/tests/models/test_qwen3_5_apply_rotary.py @@ -0,0 +1,131 @@ +import ast +import inspect + +import pytest +import torch + +from xorl.models.layers.rope import rotate_half +from xorl.models.transformers.qwen3_5 import modeling_qwen3_5 +from xorl.models.transformers.qwen3_5_moe import modeling_qwen3_5_moe +from xorl.models.transformers.qwen3_5_shared import qwen3_5_apply_rotary_pos_emb + + +pytestmark = pytest.mark.cpu + + +def _build_halved_cos_sin(batch: int, seq: int, head_dim: int) -> tuple[torch.Tensor, torch.Tensor]: + """Mimic `RotaryEmbedding.forward`: halved layout [c0..c_{d/2-1}, c0..c_{d/2-1}].""" + half = head_dim // 2 + inv_freq = 1.0 / (10000.0 ** (torch.arange(0, half, dtype=torch.float32) / half)) + positions = torch.arange(seq, dtype=torch.float32) + freqs = positions[:, None] * inv_freq[None, :] + emb = torch.cat([freqs, freqs], dim=-1) + cos = emb.cos().expand(batch, -1, -1).contiguous() + sin = emb.sin().expand(batch, -1, -1).contiguous() + return cos, sin + + +def _hf_reference_half_rotate( + q: torch.Tensor, k: torch.Tensor, cos_halved: torch.Tensor, sin_halved: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor]: + """HF/SGLang Qwen3.5 reference: standard half-rotate on q/k features.""" + cos = cos_halved.unsqueeze(2) + sin = sin_halved.unsqueeze(2) + q_embed = q * cos + rotate_half(q) * sin + k_embed = k * cos + rotate_half(k) * sin + return q_embed, k_embed + + +def _hf_reference_pairwise( + q: torch.Tensor, k: torch.Tensor, cos_halved: torch.Tensor, sin_halved: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor]: + """Pairwise-interleaved reference (DSv3 MLA decoupled-RoPE convention). + + Reshape interleaved q/k to halved, apply standard non-interleaved rotation, + then reshape back. Equivalent to `qwen3_5_apply_rotary_pos_emb(interleaved=True)`. + """ + + def to_halved(x: torch.Tensor) -> torch.Tensor: + b, s, h, d = x.shape + return x.view(b, s, h, d // 2, 2).transpose(-1, -2).reshape(b, s, h, d) + + def to_interleaved(x: torch.Tensor) -> torch.Tensor: + b, s, h, d = x.shape + return x.view(b, s, h, 2, d // 2).transpose(-1, -2).reshape(b, s, h, d) + + q_h = to_halved(q) + k_h = to_halved(k) + cos = cos_halved.unsqueeze(2) + sin = sin_halved.unsqueeze(2) + q_embed_h = q_h * cos + rotate_half(q_h) * sin + k_embed_h = k_h * cos + rotate_half(k_h) * sin + return to_interleaved(q_embed_h), to_interleaved(k_embed_h) + + +def test_default_matches_hf_half_rotate(): + """Default (interleaved=False) is the Qwen3.5/3.6 + HF/SGLang convention.""" + torch.manual_seed(0) + batch, seq, num_heads, head_dim = 2, 4, 3, 8 + q = torch.randn(batch, seq, num_heads, head_dim, dtype=torch.float32) + k = torch.randn(batch, seq, num_heads, head_dim, dtype=torch.float32) + cos, sin = _build_halved_cos_sin(batch, seq, head_dim) + + q_ours, k_ours = qwen3_5_apply_rotary_pos_emb(q, k, cos, sin) + q_ref, k_ref = _hf_reference_half_rotate(q, k, cos, sin) + + torch.testing.assert_close(q_ours, q_ref, atol=1e-6, rtol=1e-6) + torch.testing.assert_close(k_ours, k_ref, atol=1e-6, rtol=1e-6) + + +def test_interleaved_matches_pairwise_reference(): + """interleaved=True is the DSv3 MLA decoupled-RoPE pairwise convention.""" + torch.manual_seed(0) + batch, seq, num_heads, head_dim = 2, 5, 3, 8 + q = torch.randn(batch, seq, num_heads, head_dim, dtype=torch.float32) + k = torch.randn(batch, seq, num_heads, head_dim, dtype=torch.float32) + cos, sin = _build_halved_cos_sin(batch, seq, head_dim) + + q_ours, k_ours = qwen3_5_apply_rotary_pos_emb(q, k, cos, sin, interleaved=True) + q_ref, k_ref = _hf_reference_pairwise(q, k, cos, sin) + + torch.testing.assert_close(q_ours, q_ref, atol=1e-6, rtol=1e-6) + torch.testing.assert_close(k_ours, k_ref, atol=1e-6, rtol=1e-6) + + +def test_interleaved_pairwise_rotation_d8(): + """Hand-worked d=8 sanity: pair i must be rotated by angle θ_i, not θ_{i+1}.""" + torch.manual_seed(0) + batch, seq, num_heads, head_dim = 1, 1, 1, 8 + q = torch.randn(batch, seq, num_heads, head_dim, dtype=torch.float32) + k = torch.zeros_like(q) + cos, sin = _build_halved_cos_sin(batch, seq, head_dim) + + q_out, _ = qwen3_5_apply_rotary_pos_emb(q, k, cos, sin, interleaved=True) + # cos_unique[t=0] = [1,1,1,1] and sin_unique[t=0] = [0,0,0,0], so rotation + # at position 0 is the identity. + torch.testing.assert_close(q_out, q, atol=1e-6, rtol=1e-6) + + +@pytest.mark.parametrize("modeling_module", [modeling_qwen3_5, modeling_qwen3_5_moe]) +def test_qwen35_modeling_does_not_pass_interleaved_to_rotary(modeling_module): + """Regression: Qwen3.5/3.6 attention must NOT pass `interleaved=` into + `qwen3_5_apply_rotary_pos_emb`. q/k features use the standard half-rotate + convention (HF/SGLang); `mrope_interleaved` controls T/H/W frequency mixing + in cos/sin construction upstream, not the q/k rotation convention. Plumbing + it in would silently switch q/k to pairwise rotation for any HF Qwen3.5/3.6 + config that ships with `mrope_interleaved=true`. + """ + tree = ast.parse(inspect.getsource(modeling_module)) + offenders: list[int] = [] + for node in ast.walk(tree): + if ( + isinstance(node, ast.Call) + and isinstance(node.func, ast.Name) + and node.func.id == "qwen3_5_apply_rotary_pos_emb" + and any(kw.arg == "interleaved" for kw in node.keywords) + ): + offenders.append(node.lineno) + assert not offenders, ( + f"{modeling_module.__name__} calls `qwen3_5_apply_rotary_pos_emb(..., interleaved=...)` " + f"at line(s) {offenders}; this must not be plumbed from `mrope_interleaved`." + ) diff --git a/tests/models/test_qwen3_5_registry.py b/tests/models/test_qwen3_5_registry.py index bec9f4e4..d2c3d9e7 100644 --- a/tests/models/test_qwen3_5_registry.py +++ b/tests/models/test_qwen3_5_registry.py @@ -18,12 +18,18 @@ def test_qwen3_5_moe_config_from_hf_config(): "partial_rotary_factor": 0.25, "mrope_interleaved": True, } + num_hidden_layers = 40 + full_attention_interval = 4 + layer_types = [ + "full_attention" if (i + 1) % full_attention_interval == 0 else "linear_attention" + for i in range(num_hidden_layers) + ] text_config = SimpleNamespace( vocab_size=248320, hidden_size=2048, intermediate_size=2048, shared_expert_intermediate_size=512, - num_hidden_layers=40, + num_hidden_layers=num_hidden_layers, num_attention_heads=16, num_key_value_heads=2, head_dim=256, @@ -34,8 +40,8 @@ def test_qwen3_5_moe_config_from_hf_config(): use_cache=True, attention_bias=False, attention_dropout=0.0, - layer_types=["linear_attention", "full_attention"], - full_attention_interval=4, + layer_types=layer_types, + full_attention_interval=full_attention_interval, linear_num_key_heads=16, linear_num_value_heads=32, linear_key_head_dim=128, @@ -56,7 +62,7 @@ def test_qwen3_5_moe_config_from_hf_config(): config = Qwen3_5MoeConfig.from_hf_config(hf_config) - assert config.layer_types == ["linear_attention", "full_attention"] + assert config.layer_types == layer_types assert config.linear_num_key_heads == 16 assert config.linear_num_value_heads == 32 assert config.linear_key_head_dim == 128 @@ -81,6 +87,8 @@ def test_qwen3_5_config_from_hf_config(): use_cache=True, attention_bias=False, attention_dropout=0.0, + layer_types=["linear_attention", "full_attention"], + full_attention_interval=4, rope_parameters={"rope_type": "default", "rope_theta": 10_000_000}, ) hf_config = SimpleNamespace(text_config=text_config, tie_word_embeddings=False) @@ -90,3 +98,7 @@ def test_qwen3_5_config_from_hf_config(): assert config.vocab_size == 248320 assert config.hidden_size == 4096 assert config.head_dim == 256 + assert config.layer_types == [ + "full_attention" if (layer_idx + 1) % text_config.full_attention_interval == 0 else "linear_attention" + for layer_idx in range(text_config.num_hidden_layers) + ] diff --git a/tests/ops/loss/test_causallm_z_loss.py b/tests/ops/loss/test_causallm_z_loss.py new file mode 100644 index 00000000..f8e48429 --- /dev/null +++ b/tests/ops/loss/test_causallm_z_loss.py @@ -0,0 +1,172 @@ +"""Tests for the softmax auxiliary (Z-)loss term in causallm_loss_function.""" + +import pytest +import torch + +from tests.ops.loss.conftest import assert_close +from xorl.ops.loss.causallm_loss import causallm_loss_function + + +def _reference_z_loss(hidden_states, weight, labels, ignore_index=-100): + """Reference matching OLMo's Z-loss exactly. + + OLMo (olmo/train.py::cross_entropy_loss with reduction="sum", + then divided by batch_size_in_tokens in train_micro_batch) computes: + z_squared = logits.logsumexp(-1).pow(2) + z_squared = (z_squared * (labels != ignore_index)).sum() + z_loss = z_squared / num_valid_tokens + """ + h = hidden_states.view(-1, hidden_states.size(-1)) + lab = labels.view(-1) + logits = (h.float() @ weight.float().t()).float() + z_squared = logits.logsumexp(-1).pow(2) + z_squared = (z_squared * (lab != ignore_index)).sum() + return z_squared / (lab != ignore_index).sum().clamp(min=1) + + +def _reference_ce_loss(hidden_states, weight, labels, ignore_index=-100): + h = hidden_states.view(-1, hidden_states.size(-1)) + lab = labels.view(-1) + logits = (h @ weight.t()).float() + valid_count = (lab != ignore_index).sum().clamp(min=1) + per_token = torch.nn.functional.cross_entropy(logits, lab, reduction="none", ignore_index=ignore_index) + return per_token.sum() / valid_count + + +@pytest.fixture +def inputs(): + torch.manual_seed(0) + B, S, V, H = 2, 6, 32, 16 + hidden_states = torch.randn(B, S, H) / (H**0.5) + weight = torch.randn(V, H) + labels = torch.randint(0, V, (B, S)) + # Mark a couple positions as ignore so masking is exercised. + labels[0, 0] = -100 + labels[1, -1] = -100 + return hidden_states, weight, labels + + +def test_eager_z_loss_matches_reference(inputs): + hidden_states, weight, labels = inputs + ce_ref = _reference_ce_loss(hidden_states, weight, labels) + z_ref = _reference_z_loss(hidden_states, weight, labels) + coef = 1e-3 + + out = causallm_loss_function( + hidden_states=hidden_states, + weight=weight, + labels=labels, + ce_mode="eager", + z_loss_coef=coef, + ) + + assert out.metrics is not None + assert_close(out.metrics["ce_loss"], ce_ref) + assert_close(out.metrics["z_loss"], z_ref) + assert_close(out.loss, ce_ref + coef * z_ref) + + +def test_eager_no_z_loss_when_coef_zero(inputs): + hidden_states, weight, labels = inputs + ce_ref = _reference_ce_loss(hidden_states, weight, labels) + + out = causallm_loss_function( + hidden_states=hidden_states, + weight=weight, + labels=labels, + ce_mode="eager", + z_loss_coef=0.0, + ) + + assert out.metrics is None + assert_close(out.loss, ce_ref) + + +def test_eager_z_loss_grad_flows(inputs): + hidden_states, weight, labels = inputs + hidden_states = hidden_states.detach().requires_grad_(True) + weight = weight.detach().requires_grad_(True) + + out = causallm_loss_function( + hidden_states=hidden_states, + weight=weight, + labels=labels, + ce_mode="eager", + z_loss_coef=1.0, + ) + out.loss.backward() + + assert hidden_states.grad is not None and torch.isfinite(hidden_states.grad).all() + assert weight.grad is not None and torch.isfinite(weight.grad).all() + + +def test_eager_z_loss_zero_when_logits_centered(): + """If logits are all zeros, logsumexp = log(V) (constant) so Z-loss = log(V)^2.""" + torch.manual_seed(1) + B, S, V, H = 1, 3, 8, 4 + hidden_states = torch.zeros(B, S, H) + weight = torch.zeros(V, H) + labels = torch.zeros(B, S, dtype=torch.long) + + out = causallm_loss_function( + hidden_states=hidden_states, + weight=weight, + labels=labels, + ce_mode="eager", + z_loss_coef=1.0, + ) + expected_z = torch.tensor(float(torch.log(torch.tensor(V)).item() ** 2)) + assert_close(out.metrics["z_loss"], expected_z) + + +@pytest.mark.gpu +@pytest.mark.skipif(not torch.cuda.is_available(), reason="compiled CE+LSE^2 path requires CUDA") +def test_compiled_z_loss_matches_eager(inputs): + """The fused/compiled CE+LSE^2 kernel must agree with the eager reference.""" + hidden_states, weight, labels = inputs + hidden_states = hidden_states.cuda() + weight = weight.cuda() + labels = labels.cuda() + + coef = 1e-3 + out_eager = causallm_loss_function( + hidden_states=hidden_states, + weight=weight, + labels=labels, + ce_mode="eager", + z_loss_coef=coef, + ) + out_compiled = causallm_loss_function( + hidden_states=hidden_states, + weight=weight, + labels=labels, + ce_mode="compiled", + num_chunks=2, + z_loss_coef=coef, + ) + + assert_close(out_compiled.metrics["ce_loss"], out_eager.metrics["ce_loss"]) + assert_close(out_compiled.metrics["z_loss"], out_eager.metrics["z_loss"]) + assert_close(out_compiled.loss, out_eager.loss) + + +def test_tp_path_rejects_z_loss(): + """TP path must error out clearly when Z-loss is requested.""" + torch.manual_seed(2) + B, S, V, H = 1, 2, 8, 4 + hidden_states = torch.randn(B, S, H) + weight = torch.randn(V, H) + labels = torch.zeros(B, S, dtype=torch.long) + + # Pass a non-None tp_group sentinel; we want to fail before any collective. + class _Sentinel: + pass + + with pytest.raises(NotImplementedError, match="tensor parallelism"): + causallm_loss_function( + hidden_states=hidden_states, + weight=weight, + labels=labels, + tp_group=_Sentinel(), + z_loss_coef=1e-3, + ) diff --git a/tests/ops/loss/test_drgrpo_loss.py b/tests/ops/loss/test_drgrpo_loss.py index e8d0ca98..52dae725 100644 --- a/tests/ops/loss/test_drgrpo_loss.py +++ b/tests/ops/loss/test_drgrpo_loss.py @@ -10,7 +10,7 @@ import torch from tests.ops.loss.conftest import assert_close -from xorl.ops.loss import drgrpo_loss_function +from xorl.ops.loss import TokenPartial, drgrpo_loss_function @pytest.fixture @@ -105,8 +105,10 @@ def test_forward(self, inputs): 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)) + # Regression test: expected value computed with seed=42 fixture inputs. + # Default loss_reducer is TokenPartial(scale=mask.sum()); fixture has 4 + # active tokens of 8 → 2× the previous numel-scaled value. + assert_close(output.loss, torch.tensor(0.727356)) def test_backward(self, inputs): """Backward pass produces expected gradient norm (regression test).""" @@ -129,8 +131,9 @@ def test_backward(self, inputs): 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)) + # Regression test: expected value computed with seed=42 fixture inputs. + # Loss scaled 2× under new TokenPartial(scale=mask.sum()) default → grad scaled 2×. + assert_close(hidden_states.grad.norm(), torch.tensor(3.028616)) def test_zero_advantages(self, inputs): """Zero advantages produce finite (near-zero) loss.""" @@ -309,31 +312,43 @@ def test_metrics_present(self, inputs): "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.""" + def test_microbatch_composition(self, inputs): + """Per-mb partial shares sum to single-batch values for both loss and metrics. + + Regression test for the metric-inflation bug: with global-denominator + reducers, summing per-mb outputs must recover the single-batch result — + not N times as large. + """ d = inputs + B = d["B"] + assert B >= 2, "Test requires B >= 2 to form micro-batches." - losses = {} - for agg_type in ["token_mean", "fixed_horizon", "sequence_mean"]: - output = drgrpo_loss_function( - hidden_states=d["hidden_states"], + loss_mask = (d["labels_with_mask"] != d["ignore_index"]).float() + metric_reducer = TokenPartial(scale=loss_mask.sum()) + loss_reducer = TokenPartial(scale=torch.tensor(float(loss_mask.numel()))) + + def call(slc): + return drgrpo_loss_function( + hidden_states=d["hidden_states"][slc], weight=d["weight"], - labels=d["labels_with_mask"], - old_logprobs=d["old_logprobs"], - advantages=d["advantages"], + labels=d["labels_with_mask"][slc], + old_logprobs=d["old_logprobs"][slc], + advantages=d["advantages"][slc], + ref_logprobs=d["ref_logprobs"][slc], ignore_index=d["ignore_index"], - agg_type=agg_type, - beta=0.0, + beta=0.1, + loss_reducer=loss_reducer, + metric_reducer=metric_reducer, ) - 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" + single = call(slice(None)) + mbs = [call(slice(b, b + 1)) for b in range(B)] + + assert_close(sum(mb.loss for mb in mbs), single.loss) + for key, expected in single.metrics.items(): + assert_close(sum(mb.metrics[key] for mb in mbs), expected) diff --git a/tests/ops/loss/test_importance_sampling_loss.py b/tests/ops/loss/test_importance_sampling_loss.py new file mode 100644 index 00000000..8a5a1b92 --- /dev/null +++ b/tests/ops/loss/test_importance_sampling_loss.py @@ -0,0 +1,109 @@ +import pytest +import torch + +from tests.ops.loss.conftest import assert_close +from xorl.ops.loss import TokenPartial, importance_sampling_loss_function + + +_IGNORE = -100 + + +@pytest.fixture +def inputs(): + torch.manual_seed(11) + B, S, V, H = 3, 5, 12, 16 + + hidden_states = torch.randn(B, S, H) / (H**0.5) + weight = torch.randn(V, H) + labels = torch.randint(0, V, (B, S)) + + mask_pattern = torch.tensor( + [ + [1, 1, 0, 1, 0], + [1, 0, 0, 0, 0], + [1, 1, 1, 1, 1], + ], + dtype=torch.bool, + ) + labels_with_mask = labels.clone() + labels_with_mask[~mask_pattern] = _IGNORE + + return { + "B": B, + "hidden_states": hidden_states, + "weight": weight, + "labels": labels_with_mask, + "old_logprobs": torch.randn(B, S) * 0.3 - 1.5, + "advantages": torch.randn(B, S), + } + + +def _call(d, slc, *, loss_reducer=None, metric_reducer=None, **kwargs): + return importance_sampling_loss_function( + hidden_states=d["hidden_states"][slc], + weight=d["weight"], + labels=d["labels"][slc], + old_logprobs=d["old_logprobs"][slc], + advantages=d["advantages"][slc], + ignore_index=_IGNORE, + ce_mode="eager", + loss_reducer=loss_reducer, + metric_reducer=metric_reducer, + **kwargs, + ) + + +@pytest.mark.parametrize( + "extra", + [ + pytest.param({}, id="basic"), + pytest.param({"compute_kl_stats": True}, id="kl_stats"), + ], +) +def test_identity_against_legacy(inputs, extra): + d = inputs + legacy = _call(d, slice(None), **extra) + + mask = (d["labels"] != _IGNORE).float() + reducer = TokenPartial(scale=mask.sum()) + explicit = _call(d, slice(None), loss_reducer=reducer, metric_reducer=reducer, **extra) + + assert_close(explicit.loss, legacy.loss) + for key, expected in legacy.metrics.items(): + assert key in explicit.metrics + assert_close( + torch.as_tensor(explicit.metrics[key], dtype=torch.float64), + torch.as_tensor(expected, dtype=torch.float64), + ) + + +@pytest.mark.parametrize( + "extra", + [ + pytest.param({}, id="basic"), + pytest.param({"compute_kl_stats": True}, id="kl_stats"), + ], +) +def test_microbatch_composition(inputs, extra): + d = inputs + B = d["B"] + + mask = (d["labels"] != _IGNORE).float() + loss_reducer = TokenPartial(scale=mask.sum()) + metric_reducer = TokenPartial(scale=mask.sum()) + + single = _call(d, slice(None), loss_reducer=loss_reducer, metric_reducer=metric_reducer, **extra) + mbs = [ + _call(d, slice(b, b + 1), loss_reducer=loss_reducer, metric_reducer=metric_reducer, **extra) for b in range(B) + ] + + assert_close(sum(mb.loss for mb in mbs), single.loss) + + composing = {"ratio_mean", "kl_sample_train_k3", "entropy_sample"} + for key, expected in single.metrics.items(): + if key not in composing: + continue + assert_close( + torch.as_tensor(sum(mb.metrics[key] for mb in mbs), dtype=torch.float64), + torch.as_tensor(expected, dtype=torch.float64), + ) diff --git a/tests/ops/loss/test_opd_loss.py b/tests/ops/loss/test_opd_loss.py new file mode 100644 index 00000000..e3087322 --- /dev/null +++ b/tests/ops/loss/test_opd_loss.py @@ -0,0 +1,231 @@ +import pytest +import torch +from safetensors.torch import save_file + +from tests._helpers.opd import reference_opd_loss +from tests.ops.loss.conftest import assert_close +from xorl.distillation.teacher_store import TeacherHeadShardView, TeacherHeadStore, prepare_lm_head_teacher_store +from xorl.ops.loss import TokenPartial, opd_loss_function + + +pytestmark = pytest.mark.cpu + + +@pytest.fixture +def inputs(): + torch.manual_seed(7) + batch, seq, vocab, student_h, teacher_h = 2, 5, 13, 6, 8 + hidden_states = torch.randn(batch, seq, student_h) / student_h**0.5 + weight = torch.randn(vocab, student_h) / student_h**0.5 + labels = torch.randint(0, vocab, (batch, seq)) + labels[0, 0] = -100 + labels[1, -1] = -100 + teacher_hidden_states = torch.randn(batch, seq, teacher_h) / teacher_h**0.5 + teacher_weight = torch.randn(vocab, teacher_h) / teacher_h**0.5 + teacher_weights = torch.linspace(0.5, 1.5, steps=batch * seq).view(batch, seq) + return hidden_states, weight, labels, teacher_hidden_states, teacher_weight, teacher_weights + + +@pytest.mark.parametrize("num_chunks_case", [1, 2, 7, "n_valid_plus_one", 0]) +def test_opd_loss_matches_reference(inputs, num_chunks_case): + hidden_states, weight, labels, teacher_hidden_states, teacher_weight, teacher_weights = inputs + n_valid = int((labels != -100).sum().item()) + num_chunks = n_valid + 1 if num_chunks_case == "n_valid_plus_one" else int(num_chunks_case) + out = opd_loss_function( + hidden_states=hidden_states, + weight=weight, + labels=labels, + teacher_hidden_states=teacher_hidden_states, + teacher_lm_head_weight=teacher_weight, + teacher_weights=teacher_weights, + num_chunks=num_chunks, + ) + + expected = reference_opd_loss( + hidden_states, + weight, + labels, + teacher_hidden_states, + teacher_weight, + teacher_weights, + ) + assert_close(out.loss, expected) + assert out.loss.dtype == torch.float32 + assert out.metrics["valid_tokens"] == int((labels != -100).sum().item()) + + +@pytest.mark.parametrize("backend", ["streaming", "tilelang"]) +def test_opd_streaming_backends_match_reference(inputs, backend): + hidden_states, weight, labels, teacher_hidden_states, teacher_weight, teacher_weights = inputs + hidden_states = hidden_states.detach().requires_grad_(True) + weight = weight.detach().requires_grad_(True) + teacher_hidden_states = teacher_hidden_states.detach().requires_grad_(True) + teacher_weight = teacher_weight.detach().requires_grad_(True) + + out = opd_loss_function( + hidden_states=hidden_states, + weight=weight, + labels=labels, + teacher_hidden_states=teacher_hidden_states, + teacher_lm_head_weight=teacher_weight, + teacher_weights=teacher_weights, + kl_backend=backend, + vocab_chunk_size=5, + ) + + expected = reference_opd_loss( + hidden_states, + weight, + labels, + teacher_hidden_states, + teacher_weight, + teacher_weights, + ) + assert_close(out.loss, expected) + out.loss.backward() + assert hidden_states.grad is not None and hidden_states.grad.isfinite().all() + assert weight.grad is not None and weight.grad.isfinite().all() + assert teacher_hidden_states.grad is None + assert teacher_weight.grad is None + + +def test_opd_streaming_backend_reads_sharded_teacher_store(inputs, tmp_path): + hidden_states, weight, labels, teacher_hidden_states, teacher_weight, teacher_weights = inputs + model_dir = tmp_path / "teacher_model" + model_dir.mkdir() + save_file({"lm_head.weight": teacher_weight}, str(model_dir / "model.safetensors")) + manifest = prepare_lm_head_teacher_store(model_dir, tmp_path / "teacher_store", teacher_id=0, shard_rows=4) + teacher_view = TeacherHeadShardView( + store=TeacherHeadStore(manifest), + teacher_id="0", + device=torch.device("cpu"), + dtype=None, + ) + + out = opd_loss_function( + hidden_states=hidden_states, + weight=weight, + labels=labels, + teacher_hidden_states=teacher_hidden_states, + teacher_lm_head_weight=teacher_view, + teacher_weights=teacher_weights, + kl_backend="streaming", + vocab_chunk_size=5, + ) + + expected = reference_opd_loss(hidden_states, weight, labels, teacher_hidden_states, teacher_weight, teacher_weights) + assert_close(out.loss, expected) + + +def test_opd_loss_backward(inputs): + hidden_states, weight, labels, teacher_hidden_states, teacher_weight, _ = inputs + hidden_states = hidden_states.detach().requires_grad_(True) + weight = weight.detach().requires_grad_(True) + teacher_hidden_states = teacher_hidden_states.detach().requires_grad_(True) + teacher_weight = teacher_weight.detach().requires_grad_(True) + + out = opd_loss_function( + hidden_states=hidden_states, + weight=weight, + labels=labels, + teacher_hidden_states=teacher_hidden_states, + teacher_lm_head_weight=teacher_weight, + num_chunks=2, + ) + out.loss.backward() + + assert hidden_states.grad is not None and hidden_states.grad.isfinite().all() + assert weight.grad is not None and weight.grad.isfinite().all() + assert teacher_hidden_states.grad is None + assert teacher_weight.grad is None + + +def test_opd_loss_respects_token_partial_reducer(inputs): + hidden_states, weight, labels, teacher_hidden_states, teacher_weight, teacher_weights = inputs + out = opd_loss_function( + hidden_states=hidden_states, + weight=weight, + labels=labels, + teacher_hidden_states=teacher_hidden_states, + teacher_lm_head_weight=teacher_weight, + teacher_weights=teacher_weights, + num_chunks=2, + loss_reducer=TokenPartial(scale=torch.tensor(1.0)), + ) + + n_valid = (labels != -100).sum().to(dtype=torch.float32) + expected = reference_opd_loss(hidden_states, weight, labels, teacher_hidden_states, teacher_weight, teacher_weights) + assert_close(out.loss, expected * n_valid) + + +def test_opd_loss_all_ignored_is_finite(inputs): + hidden_states, weight, labels, teacher_hidden_states, teacher_weight, _ = inputs + labels = torch.full_like(labels, -100) + hidden_states = hidden_states.to(torch.bfloat16).detach().requires_grad_(True) + weight = weight.to(torch.bfloat16).detach().requires_grad_(True) + + out = opd_loss_function( + hidden_states=hidden_states, + weight=weight, + labels=labels, + teacher_hidden_states=teacher_hidden_states, + teacher_lm_head_weight=teacher_weight, + lm_head_fp32=True, + ) + + assert out.loss.isfinite() + assert out.loss.item() == 0.0 + assert out.loss.dtype == torch.float32 + assert out.metrics["valid_tokens"] == 0 + out.loss.backward() + assert hidden_states.grad is not None and hidden_states.grad.isfinite().all() + assert weight.grad is not None and weight.grad.isfinite().all() + assert torch.count_nonzero(hidden_states.grad) == 0 + assert torch.count_nonzero(weight.grad) == 0 + + +def test_opd_loss_bf16_inputs_return_fp32_loss(inputs): + hidden_states, weight, labels, teacher_hidden_states, teacher_weight, teacher_weights = inputs + hidden_states = hidden_states.to(torch.bfloat16).detach().requires_grad_(True) + weight = weight.to(torch.bfloat16).detach().requires_grad_(True) + teacher_hidden_states = teacher_hidden_states.to(torch.bfloat16) + teacher_weight = teacher_weight.to(torch.bfloat16) + + out = opd_loss_function( + hidden_states=hidden_states, + weight=weight, + labels=labels, + teacher_hidden_states=teacher_hidden_states, + teacher_lm_head_weight=teacher_weight, + teacher_weights=teacher_weights, + num_chunks=2, + lm_head_fp32=True, + teacher_lm_head_fp32=True, + ) + + assert out.loss.dtype == torch.float32 + out.loss.backward() + assert hidden_states.grad is not None and hidden_states.grad.isfinite().all() + assert weight.grad is not None and weight.grad.isfinite().all() + + +def test_opd_loss_return_per_token(inputs): + hidden_states, weight, labels, teacher_hidden_states, teacher_weight, teacher_weights = inputs + out = opd_loss_function( + hidden_states=hidden_states, + weight=weight, + labels=labels, + teacher_hidden_states=teacher_hidden_states, + teacher_lm_head_weight=teacher_weight, + teacher_weights=teacher_weights, + num_chunks=2, + return_per_token=True, + ) + + assert out.per_token_loss is not None + assert out.per_token_loss.shape == labels.shape + assert out.per_token_loss.dtype == torch.float32 + assert torch.count_nonzero(out.per_token_loss[labels == -100]) == 0 + expected = reference_opd_loss(hidden_states, weight, labels, teacher_hidden_states, teacher_weight, teacher_weights) + denom = (labels != -100).sum().to(dtype=torch.float32) + assert_close(out.per_token_loss.sum() / denom, expected) diff --git a/tests/ops/loss/test_policy_loss.py b/tests/ops/loss/test_policy_loss.py new file mode 100644 index 00000000..8bf36962 --- /dev/null +++ b/tests/ops/loss/test_policy_loss.py @@ -0,0 +1,129 @@ +import pytest +import torch + +from tests.ops.loss.conftest import assert_close +from xorl.ops.loss import TokenPartial, policy_loss_function + + +_IGNORE = -100 + + +@pytest.fixture +def inputs(): + torch.manual_seed(7) + B, S, V, H = 3, 5, 12, 16 + + hidden_states = torch.randn(B, S, H) / (H**0.5) + weight = torch.randn(V, H) + labels = torch.randint(0, V, (B, S)) + + # Non-uniform mask across rows so each mb has a different valid count. + mask_pattern = torch.tensor( + [ + [1, 1, 0, 1, 0], + [1, 0, 0, 0, 0], + [1, 1, 1, 1, 1], + ], + dtype=torch.bool, + ) + labels_with_mask = labels.clone() + labels_with_mask[~mask_pattern] = _IGNORE + + return { + "B": B, + "hidden_states": hidden_states, + "weight": weight, + "labels": labels_with_mask, + "old_logprobs": torch.randn(B, S) * 0.3 - 1.5, + "rollout_logprobs": torch.randn(B, S) * 0.3 - 1.5, + "advantages": torch.randn(B, S), + } + + +def _call(d, slc, *, loss_reducer=None, metric_reducer=None, **kwargs): + return policy_loss_function( + hidden_states=d["hidden_states"][slc], + weight=d["weight"], + labels=d["labels"][slc], + old_logprobs=d["old_logprobs"][slc], + advantages=d["advantages"][slc], + rollout_logprobs=d["rollout_logprobs"][slc], + ignore_index=_IGNORE, + ce_mode="eager", + loss_reducer=loss_reducer, + metric_reducer=metric_reducer, + **kwargs, + ) + + +@pytest.mark.parametrize( + "extra", + [ + pytest.param({}, id="vanilla_ppo"), + pytest.param({"use_tis": True}, id="tis"), + pytest.param({"icepop_beta": 1.5}, id="icepop"), + pytest.param({"compute_kl_stats": True}, id="kl_stats"), + ], +) +def test_identity_against_legacy(inputs, extra): + """``TokenPartial(scale=mask.sum())`` reproduces the legacy local-mean result.""" + d = inputs + legacy = _call(d, slice(None), **extra) + + mask = (d["labels"] != _IGNORE).float() + reducer = TokenPartial(scale=mask.sum()) + explicit = _call(d, slice(None), loss_reducer=reducer, metric_reducer=reducer, **extra) + + assert_close(explicit.loss, legacy.loss) + for key, expected in legacy.metrics.items(): + assert key in explicit.metrics + # Ratio min/max are local reductions — identical regardless of reducer. + # Mean metrics match because TokenPartial(scale=mask.sum()) ≡ legacy mean. + assert_close( + torch.as_tensor(explicit.metrics[key], dtype=torch.float64), + torch.as_tensor(expected, dtype=torch.float64), + ) + + +@pytest.mark.parametrize( + "extra", + [ + pytest.param({}, id="vanilla_ppo"), + pytest.param({"use_tis": True}, id="tis"), + pytest.param({"icepop_beta": 1.5}, id="icepop"), + pytest.param({"compute_kl_stats": True}, id="kl_stats"), + ], +) +def test_microbatch_composition(inputs, extra): + """Sum of partial shares across non-uniform mbs equals the single-batch value.""" + d = inputs + B = d["B"] + + mask = (d["labels"] != _IGNORE).float() + loss_reducer = TokenPartial(scale=mask.sum()) + metric_reducer = TokenPartial(scale=mask.sum()) + + single = _call(d, slice(None), loss_reducer=loss_reducer, metric_reducer=metric_reducer, **extra) + mbs = [ + _call(d, slice(b, b + 1), loss_reducer=loss_reducer, metric_reducer=metric_reducer, **extra) for b in range(B) + ] + + assert_close(sum(mb.loss for mb in mbs), single.loss) + + # Reducer-routed metrics should compose under sum. + composing = { + "pg_clipfrac", + "icepop_maskfrac", + "tis_mean", + "tis_clipfrac", + "kl_sample_train_k3", + "entropy_sample", + "ratio_mean", + } + for key, expected in single.metrics.items(): + if key not in composing: + continue + assert_close( + torch.as_tensor(sum(mb.metrics[key] for mb in mbs), dtype=torch.float64), + torch.as_tensor(expected, dtype=torch.float64), + ) diff --git a/tests/ops/loss/test_reducers.py b/tests/ops/loss/test_reducers.py new file mode 100644 index 00000000..77028fa0 --- /dev/null +++ b/tests/ops/loss/test_reducers.py @@ -0,0 +1,197 @@ +"""Direct tests for the Reducer abstraction. + +The contract: a reducer is a closure over a caller-supplied denominator, so its +outputs are partial shares — they sum across micro-batches (and across ranks +under all_reduce(SUM)) to the globally-correct value. +""" + +import pytest +import torch + +from tests.ops.loss.conftest import assert_close +from xorl.ops.loss import SequencePartial, TokenPartial + + +@pytest.mark.parametrize( + "scale_fn", + [ + pytest.param(lambda mask: mask.sum(), id="active_count"), + pytest.param(lambda mask: torch.tensor(float(mask.numel())), id="numel"), + pytest.param(lambda mask: torch.tensor(1.0), id="ones_raw_sum"), + ], +) +def test_token_partial_shares_sum(scale_fn): + """Per-microbatch TokenPartial shares sum to the single-batch result.""" + torch.manual_seed(0) + B, S = 4, 6 + values = torch.randn(B, S) + mask = torch.randint(0, 2, (B, S)).float() + + reducer = TokenPartial(scale=scale_fn(mask)) + single = reducer(values, mask) + summed = sum(reducer(values[b : b + 1], mask[b : b + 1]) for b in range(B)) + + assert_close(summed, single) + + +def test_token_partial_scale_one_equals_raw_sum(): + """``TokenPartial(scale=1)`` is the raw masked sum (deferred-divide form).""" + torch.manual_seed(0) + B, S = 4, 6 + values = torch.randn(B, S) + mask = torch.randint(0, 2, (B, S)).float() + + out = TokenPartial(scale=torch.tensor(1.0))(values, mask) + assert_close(out, (values * mask).sum()) + + +def test_sequence_partial_shares_sum(): + """Per-microbatch SequencePartial shares (with sliced cu_seqlens_local) sum to the single-batch result.""" + torch.manual_seed(0) + B, S = 4, 6 + values = torch.randn(B, S) + mask = torch.randint(0, 2, (B, S)).float() + seq_lengths = mask.sum(dim=-1) + seq_count = torch.tensor(float(B)) + full_cu_seqlens = torch.arange(0, B * S + 1, S) + mb_cu_seqlens = torch.arange(0, S + 1, S) + + single = SequencePartial( + scale=seq_count, + cu_seqlens_local=full_cu_seqlens, + seq_lengths_global=seq_lengths, + )(values, mask) + summed = sum( + SequencePartial( + scale=seq_count, + cu_seqlens_local=mb_cu_seqlens, + seq_lengths_global=seq_lengths[b : b + 1], + )(values[b : b + 1], mask[b : b + 1]) + for b in range(B) + ) + + assert_close(summed, single) + + +def test_sequence_partial_scale_one_equals_sum_of_per_seq_means(): + """``SequencePartial(scale=1)`` is the deferred-outer-divide form of SequencePartial(scale=n_seqs).""" + torch.manual_seed(0) + B, S = 4, 6 + values = torch.randn(B, S) + mask = torch.randint(0, 2, (B, S)).float() + seq_lengths = mask.sum(dim=-1) + seq_count = torch.tensor(float(B)) + cu_seqlens_local = torch.arange(0, B * S + 1, S) + + deferred = SequencePartial( + scale=torch.tensor(1.0), + cu_seqlens_local=cu_seqlens_local, + seq_lengths_global=seq_lengths, + )(values, mask) + finalized = SequencePartial( + scale=seq_count, + cu_seqlens_local=cu_seqlens_local, + seq_lengths_global=seq_lengths, + )(values, mask) + + assert_close(deferred / seq_count, finalized) + + +@pytest.mark.parametrize( + "reducer", + [ + pytest.param(TokenPartial(scale=torch.tensor(0.0)), id="token"), + pytest.param( + SequencePartial( + scale=torch.tensor(0.0), + cu_seqlens_local=torch.tensor([0, 4, 8]), + seq_lengths_global=torch.zeros(2), + ), + id="sequence", + ), + ], +) +def test_empty_mask_yields_zero(reducer): + """Zero denominators clamp to 1 and produce 0, not NaN.""" + values = torch.randn(2, 4) + mask = torch.zeros(2, 4) + assert reducer(values, mask) == 0.0 + + +def test_sequence_partial_packed_row_matches_per_segment_mean(): + """One row, three packed segments described by ``cu_seqlens`` — sum of per-segment means / n_seqs.""" + torch.manual_seed(0) + # Row of 10 tokens packing three segments of lengths [3, 4, 3]. + values = torch.randn(1, 10) + mask = torch.ones(1, 10) + cu_seqlens = torch.tensor([0, 3, 7, 10]) + seg_lengths = torch.tensor([3, 4, 3]) + n_seqs = torch.tensor(float(seg_lengths.numel())) + + out = SequencePartial( + scale=n_seqs, + cu_seqlens_local=cu_seqlens, + seq_lengths_global=seg_lengths, + )(values, mask) + + flat = (values * mask).flatten() + expected = (flat[0:3].sum() / 3 + flat[3:7].sum() / 4 + flat[7:10].sum() / 3) / n_seqs + + assert_close(out, expected) + + +def test_sequence_partial_packed_cp_shares_sum(): + """Packed row split across two CP shards: per-shard partials sum to the single-batch result. + + Layout: one row of 10 tokens with three packed segments of pre-shard lengths + [3, 4, 3]. CP=2 splits the row at column 5: + + - Shard 0 covers columns [0, 5): segment 0 lives wholly here ([0, 3)), + and the first half of segment 1 ([3, 5), local length 2 of the + pre-shard 4). + - Shard 1 covers columns [5, 10): the second half of segment 1 + ([5, 7), local length 2 of 4) and segment 2 wholly ([7, 10)). + + Both shards reference the same ``seq_lengths_global`` for any segment + they touch. Segment 1's contributions from the two shards each divide + by 4, then sum to the correct full-segment mean. + """ + torch.manual_seed(0) + values = torch.randn(1, 10) + mask = torch.ones(1, 10) + seg_lengths_global = torch.tensor([3, 4, 3]) + n_seqs = torch.tensor(float(seg_lengths_global.numel())) + + full = SequencePartial( + scale=n_seqs, + cu_seqlens_local=torch.tensor([0, 3, 7, 10]), + seq_lengths_global=seg_lengths_global, + )(values, mask) + + shard_0 = SequencePartial( + scale=n_seqs, + cu_seqlens_local=torch.tensor([0, 3, 5]), + seq_lengths_global=seg_lengths_global[:2], + )(values[:, :5], mask[:, :5]) + + shard_1 = SequencePartial( + scale=n_seqs, + cu_seqlens_local=torch.tensor([0, 2, 5]), + seq_lengths_global=seg_lengths_global[1:], + )(values[:, 5:], mask[:, 5:]) + + assert_close(shard_0 + shard_1, full) + + +def test_token_partial_with_n_seqs_scale_equals_seq_mean_token_sum(): + """``TokenPartial(scale=n_seqs)`` expresses verl's seq-mean-token-sum policy.""" + torch.manual_seed(0) + B, S = 4, 6 + values = torch.randn(B, S) + mask = torch.randint(0, 2, (B, S)).float() + seq_count = torch.tensor(float(B)) + + via_token_partial = TokenPartial(scale=seq_count)(values, mask) + direct = (values * mask).sum(dim=-1).sum() / seq_count + + assert_close(via_token_partial, direct) diff --git a/tests/ops/test_attention.py b/tests/ops/test_attention.py index 0ab07c5d..5bc531fb 100644 --- a/tests/ops/test_attention.py +++ b/tests/ops/test_attention.py @@ -69,6 +69,8 @@ def test_flash_attention_api_behavior(self): """Warnings, is_causal handling, return values, scaling, sliding window.""" module = Mock() module.is_causal = True + module.config = Mock() + module.config._flash_attention_deterministic = False batch, seqlen, num_heads, head_dim = 2, 16, 8, 64 query = torch.randn(batch, seqlen, num_heads, head_dim) @@ -148,10 +150,23 @@ def test_flash_attention_api_behavior(self): ) assert mock_fa.call_args[1]["window_size"] == (128, 0) + # Configured deterministic backward flag is forwarded to FA3. + module.config._flash_attention_deterministic = True + flash_attention_forward( + module, + query, + key, + value, + attention_mask=None, + ) + assert mock_fa.call_args[1]["deterministic"] is True + def test_varlen_path_with_cu_seqlens(self): """cu_seqlens kwargs trigger the varlen path.""" module = Mock() module.is_causal = True + module.config = Mock() + module.config._flash_attention_deterministic = True total_tokens, num_heads, head_dim = 32, 8, 64 query = torch.randn(1, total_tokens, num_heads, head_dim) @@ -174,6 +189,7 @@ def test_varlen_path_with_cu_seqlens(self): ) assert mock_varlen.called assert mock_varlen.call_args[1]["cu_seqlens_q"].dtype == torch.int32 + assert mock_varlen.call_args[1]["deterministic"] is True assert result.shape == (1, total_tokens, num_heads, head_dim) diff --git a/tests/ops/test_ep_adapter_wrappers.py b/tests/ops/test_ep_adapter_wrappers.py index 94e90053..0af2edf1 100644 --- a/tests/ops/test_ep_adapter_wrappers.py +++ b/tests/ops/test_ep_adapter_wrappers.py @@ -7,6 +7,7 @@ """ import importlib.util +import inspect import sys import types from pathlib import Path @@ -15,6 +16,7 @@ import torch import torch.nn.functional as F +from xorl.models.layers.moe.backend import EP_EXPERT_COMPUTE from xorl.utils import import_utils @@ -159,23 +161,37 @@ def test_adapter_forwards_expert_scores(monkeypatch, backend_type, class_name): ) = _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) + output = fn_cls.apply(permute_tokens, cumsum, gate_up_proj, down_proj, intermediate_size, 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. +# --------------------------------------------------------------------------- +# Signature contract tests — replaces the old source-grep regression test. +# Inspects live function signatures instead of pattern-matching on source text, +# so formatting changes, variable renames, and line rewrapping can't fool it. +# --------------------------------------------------------------------------- - 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() +# Required explicit parameters for every EP_EXPERT_COMPUTE entry. +_REQUIRED_EP_PARAMS = ("expert_scores", "hidden_act") - 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()" + +@pytest.mark.parametrize("name,fn", list(EP_EXPERT_COMPUTE.items())) +def test_ep_compute_signature_contract(name, fn): + """Every EP compute function must explicitly accept expert_scores, hidden_act, + and a **kwargs catch-all for forward-compat extras (gate_up_bias, down_bias, etc.). + """ + sig = inspect.signature(fn) + params = sig.parameters + + for required in _REQUIRED_EP_PARAMS: + assert required in params, ( + f"EP_EXPERT_COMPUTE['{name}'] is missing explicit '{required}' param. Signature: {sig}" + ) + + has_var_keyword = any(p.kind == p.VAR_KEYWORD for p in params.values()) + assert has_var_keyword, ( + f"EP_EXPERT_COMPUTE['{name}'] has no **kwargs — new extras like " + f"gate_up_bias will break callers. Signature: {sig}" ) diff --git a/tests/ops/test_ep_routing_scores.py b/tests/ops/test_ep_routing_scores.py index a4a1b6ad..4ce7e7b9 100644 --- a/tests/ops/test_ep_routing_scores.py +++ b/tests/ops/test_ep_routing_scores.py @@ -63,6 +63,7 @@ def _patch_ep_kernels(monkeypatch, module_name: str): moe_stub.moe_gather = None moe_stub.moe_index_compute = None moe_stub.moe_scatter = None + moe_stub.moe_add_gather = 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) @@ -152,27 +153,16 @@ def test_ep_group_gemm_propagates_routing_score_gradients(monkeypatch, module_na expert_scores = torch.rand(num_tokens, dtype=dtype, requires_grad=True) upstream = torch.randn(num_tokens, hidden_dim, dtype=dtype) - # 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, - ) + # Both TritonEPGroupGemm and QuackEPGroupGemm take a fused gate_up_proj + intermediate_size (int). + 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, + ) output.backward(upstream) grad_scores = expert_scores.grad.detach().clone() diff --git a/tests/ops/test_gated_delta_rule.py b/tests/ops/test_gated_delta_rule.py index 2a1b3abd..0386c5db 100644 --- a/tests/ops/test_gated_delta_rule.py +++ b/tests/ops/test_gated_delta_rule.py @@ -1,8 +1,7 @@ import pytest import torch import torch.nn.functional as F - -from xorl.ops.linear_attention.ops.gated_delta_rule import chunk_gated_delta_rule +from fla.ops.gated_delta_rule import chunk_gated_delta_rule pytestmark = [ diff --git a/tests/ops/test_lora_utils.py b/tests/ops/test_lora_utils.py new file mode 100644 index 00000000..0c34ca72 --- /dev/null +++ b/tests/ops/test_lora_utils.py @@ -0,0 +1,41 @@ +"""Tests for stacked LoRA helper utilities.""" + +import pytest +import torch + +from xorl.ops.group_gemm.kernel.lora_utils import ( + get_lora_delta_weight_stacked, + init_lora_weights_stacked, + merge_lora_weights_stacked, + unmerge_lora_weights_stacked, +) + + +pytestmark = [pytest.mark.cpu] + + +def test_stacked_lora_helpers_use_gkn_layout(): + lora_A, lora_B = init_lora_weights_stacked( + num_experts=2, + r=3, + in_features=4, + out_features=5, + ) + + assert lora_A.shape == (2, 4, 3) + assert lora_B.shape == (2, 3, 5) + assert torch.equal(lora_B, torch.zeros_like(lora_B)) + + lora_A = torch.arange(2 * 4 * 3, dtype=torch.float32).reshape(2, 4, 3) + lora_B = torch.arange(2 * 3 * 5, dtype=torch.float32).reshape(2, 3, 5) + base = torch.ones(2, 4, 5, dtype=torch.float32) + scaling = 0.25 + expected_delta = torch.bmm(lora_A, lora_B) * scaling + + delta = get_lora_delta_weight_stacked(lora_A, lora_B, scaling) + merged = merge_lora_weights_stacked(base, lora_A, lora_B, scaling) + unmerged = unmerge_lora_weights_stacked(merged, lora_A, lora_B, scaling) + + assert torch.equal(delta, expected_delta) + assert torch.equal(merged, base + expected_delta) + assert torch.equal(unmerged, base) diff --git a/tests/ops/test_moe_act.py b/tests/ops/test_moe_act.py deleted file mode 100644 index 4616740d..00000000 --- a/tests/ops/test_moe_act.py +++ /dev/null @@ -1,398 +0,0 @@ -"""Tests for moe_act selective activation recompute variants. - -Each backend (native, triton, quack) has a moe_act variant that: -- Checkpoints gate+up projection activations (native: torch.utils.checkpoint; - triton/quack: custom autograd.Function saves fewer tensors) -- Recomputes gate+up in backward instead of saving them -- Trades extra backward compute for reduced activation memory - -Tests: -1. Forward correctness: moe_act output matches standard output per backend -2. Backward correctness: moe_act gradients match standard gradients per backend -3. Memory: moe_act saves activation memory compared to standard -4. TFLOPS benchmark: standard vs moe_act fwd+bwd per backend -""" - -import pytest -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 - - -# --------------------------------------------------------------------------- -# Helpers -# --------------------------------------------------------------------------- - - -def _available_moe_act_backends(): - """Return backends that have moe_act local variants registered.""" - 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 [] - - -def _make_block_pair(ne, hd, inter, topk, backend, seed=42): - """Create (standard, moe_act) MoEBlock pair with identical weights. - - Returns (std_block, act_block) where act_block.experts._moe_act = True. - """ - from xorl.models.layers.moe.moe_block import MoEBlock # noqa: PLC0415 - - torch.manual_seed(seed) - std = MoEBlock(hd, ne, topk, inter, moe_implementation=backend) - nn.init.xavier_normal_(std.experts.gate_proj.data) - nn.init.xavier_normal_(std.experts.up_proj.data) - nn.init.xavier_normal_(std.experts.down_proj.data) - nn.init.xavier_normal_(std.gate.weight.data) - std = std.to(DEVICE, DTYPE) - - act = MoEBlock(hd, ne, topk, inter, moe_implementation=backend) - act.experts._moe_act = True - act = act.to(DEVICE, DTYPE) - with torch.no_grad(): - act.gate.weight.copy_(std.gate.weight) - act.experts.gate_proj.copy_(std.experts.gate_proj) - act.experts.up_proj.copy_(std.experts.up_proj) - act.experts.down_proj.copy_(std.experts.down_proj) - - return std, act - - -# --------------------------------------------------------------------------- -# Test 1: Forward correctness -# --------------------------------------------------------------------------- - -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 -] - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize("backend", AVAILABLE_BACKENDS) -@pytest.mark.parametrize("ne,hd,inter,topk,bs,seq", CORRECTNESS_CONFIGS) -def test_forward_correctness(backend, ne, hd, inter, topk, bs, seq): - """moe_act forward output must match standard forward output.""" - std, act = _make_block_pair(ne, hd, inter, topk, backend) - - torch.manual_seed(7) - x = torch.randn(bs, seq, hd, device=DEVICE, dtype=DTYPE) - - with torch.no_grad(): - std_out, std_logits = std(x) - act_out, act_logits = act(x) - - # Router logits must be identical (same gate weights, same input) - torch.testing.assert_close(act_logits, std_logits, atol=0, rtol=0) - - max_diff = (act_out - std_out).abs().max().item() - torch.testing.assert_close( - act_out, - std_out, - atol=0.05, - rtol=0.02, - msg=f"[{backend}] Forward mismatch: max_diff={max_diff:.6f}", - ) - - -# --------------------------------------------------------------------------- -# Test 2: Backward correctness (gradients) -# --------------------------------------------------------------------------- - -BACKWARD_CONFIGS = [ - (4, 64, 128, 2, 2, 8), - (8, 128, 256, 2, 4, 16), -] - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize("backend", AVAILABLE_BACKENDS) -@pytest.mark.parametrize("ne,hd,inter,topk,bs,seq", BACKWARD_CONFIGS) -def test_backward_correctness(backend, ne, hd, inter, topk, bs, seq): - """moe_act backward gradients must match standard backward gradients.""" - std, act = _make_block_pair(ne, hd, inter, topk, backend) - - atol, rtol = 0.05, 0.05 - - torch.manual_seed(7) - x_std = torch.randn(bs, seq, hd, device=DEVICE, dtype=DTYPE, requires_grad=True) - x_act = x_std.detach().clone().requires_grad_(True) - - std_out, _ = std(x_std) - std_out.sum().backward() - - act_out, _ = act(x_act) - act_out.sum().backward() - - torch.testing.assert_close( - 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"]: - g_std = getattr(std.experts, name).grad - g_act = getattr(act.experts, name).grad - 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, - msg=f"[{backend}] {name} gradient mismatch", - ) - torch.testing.assert_close( - act.gate.weight.grad, - std.gate.weight.grad, - atol=atol, - rtol=rtol, - msg=f"[{backend}] Gate weight gradient mismatch", - ) - - -# --------------------------------------------------------------------------- -# 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): - out, _ = block(x) - out.sum().backward() - block.zero_grad() - if x.grad is not None: - x.grad = None - - torch.cuda.synchronize() - torch.cuda.reset_peak_memory_stats() - torch.cuda.empty_cache() - - x_clone = x.detach().clone().requires_grad_(True) - out, _ = block(x_clone) - out.sum().backward() - - torch.cuda.synchronize() - peak = torch.cuda.max_memory_allocated() - block.zero_grad() - return peak - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize("backend", AVAILABLE_BACKENDS) -def test_memory_savings(backend): - """moe_act should use less or equal peak memory than standard.""" - ne, hd, inter, topk, bs, seq = 8, 512, 1024, 2, 4, 256 - - std, act = _make_block_pair(ne, hd, inter, topk, backend) - - x = torch.randn(bs, seq, hd, device=DEVICE, dtype=DTYPE, requires_grad=True) - - mem_std = _measure_fwd_bwd_peak_memory(std, x) - mem_act = _measure_fwd_bwd_peak_memory(act, x) - - 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"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: {mem_act / 1024**2:.1f} MB vs {mem_std / 1024**2:.1f} MB" - ) - - -# --------------------------------------------------------------------------- -# 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 - - -def _benchmark(block, x, warmup=20, iters=40): - """Median fwd+bwd GPU time (seconds).""" - for _ in range(warmup): - out, _ = block(x) - out.sum().backward() - block.zero_grad() - if x.grad is not None: - x.grad = None - - torch.cuda.synchronize() - times = [] - for _ in range(iters): - t0 = torch.cuda.Event(enable_timing=True) - t1 = torch.cuda.Event(enable_timing=True) - t0.record() - out, _ = block(x) - out.sum().backward() - t1.record() - torch.cuda.synchronize() - times.append(t0.elapsed_time(t1) / 1000.0) - block.zero_grad() - if x.grad is not None: - x.grad = None - - times.sort() - return times[len(times) // 2] - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize("seq_len", [1024, 4096]) -def bench_moe_act_tflops(seq_len): - """TFLOPS benchmark: standard vs moe_act per backend.""" - ne, hd, inter, topk, bs = 8, 1024, 2048, 2, 4 - - flops = _moe_flops(bs, seq_len, hd, inter, topk) - x = torch.randn(bs, seq_len, hd, device=DEVICE, dtype=DTYPE, requires_grad=True) - - results = {} - for backend in AVAILABLE_BACKENDS: - std, act = _make_block_pair(ne, hd, inter, topk, backend) - - t_std = _benchmark(std, x) - t_act = _benchmark(act, x) - - results[backend] = { - "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, - } - - del std, act - - print("\n" + "=" * 90) - 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} {'Std TF':>10} {'Act TF':>10} {'Overhead':>10}") - print("-" * 90) - for backend, r in results.items(): - print( - f" {backend:<10} {r['std_ms']:>10.2f} {r['act_ms']:>10.2f} " - f"{r['std_tflops']:>10.2f} {r['act_tflops']:>10.2f} " - f"{r['overhead']:>9.2f}x" - ) - print("=" * 90) - - -# --------------------------------------------------------------------------- -# Test 4: moe_act works correctly inside gradient_checkpointing_enable -# --------------------------------------------------------------------------- - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize("backend", AVAILABLE_BACKENDS) -def test_moe_act_via_gradient_checkpointing_enable(backend): - """gradient_checkpointing_enable(moe_checkpoint_method='moe_act') sets _moe_act correctly.""" - from xorl.models.layers.moe.experts import MoEExperts # noqa: PLC0415 - - config = Qwen3MoeConfig( - vocab_size=1000, - num_hidden_layers=2, - hidden_size=128, - intermediate_size=256, - num_attention_heads=4, - num_key_value_heads=2, - moe_intermediate_size=128, - num_experts=4, - num_experts_per_tok=2, - decoder_sparse_step=1, - norm_topk_prob=True, - output_router_logits=False, - _moe_implementation=backend, - max_position_embeddings=128, - pad_token_id=0, - _attn_implementation="sdpa", - ) - 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", - } - ) - - # Verify _moe_act is set on all MoEExperts modules - moe_experts_modules = [m for m in model.modules() if isinstance(m, MoEExperts)] - assert len(moe_experts_modules) > 0, "No MoEExperts found in model" - for mod in moe_experts_modules: - assert mod._moe_act is True, f"Expected _moe_act=True, got {mod._moe_act}" - - # Forward + backward must not crash - input_ids = torch.randint(0, 1000, (2, 16), device=DEVICE) - output = model(input_ids=input_ids) - output.last_hidden_state.sum().backward() - - has_grad = any(p.grad is not None for p in model.parameters() if p.requires_grad) - assert has_grad, "No gradients computed" - - -# --------------------------------------------------------------------------- -# 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): - """moe_act variant must produce correct results when compiled with inductor.""" - ne, hd, inter, topk = 4, 128, 256, 2 - _, act = _make_block_pair(ne, hd, inter, topk, backend) - - # Reference: uncompiled moe_act - torch.manual_seed(42) - x = torch.randn(2, 16, hd, device=DEVICE, dtype=DTYPE, requires_grad=True) - ref_out, _ = act(x) - ref_out.sum().backward() - ref_grad = x.grad.clone() - act.zero_grad() - - # Compiled variant - torch._dynamo.reset() - _, act2 = _make_block_pair(ne, hd, inter, topk, backend) - compiled = torch.compile(act2, fullgraph=False, backend="inductor", dynamic=False) - - x2 = x.detach().clone().requires_grad_(True) - 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._dynamo.reset() - - -if __name__ == "__main__": - pytest.main([__file__, "-v", "-s"]) diff --git a/tests/ops/test_moe_gkn_format.py b/tests/ops/test_moe_gkn_format.py index 8e4e547d..dc972111 100644 --- a/tests/ops/test_moe_gkn_format.py +++ b/tests/ops/test_moe_gkn_format.py @@ -222,7 +222,7 @@ def test_backends_match_reference_and_agree(self): gate_gkn, up_gkn, down_gkn = make_gkn_weights_from_linear(experts) hidden_states_cpu = torch.randn(num_tokens, hidden_size) - act_fn = torch.nn.SiLU() + act_fn = "silu" for expert_idx in range(num_experts): ref_out = experts[expert_idx](hidden_states_cpu) @@ -260,6 +260,7 @@ def test_backends_match_reference_and_agree(self): raise ImportError from xorl.ops.moe.triton import TritonMoeExpertsFunction # noqa: PLC0415 + gate_up_cuda = torch.cat([gate_cuda, up_cuda], dim=-1) triton_out = TritonMoeExpertsFunction.apply( num_experts, rw, @@ -268,6 +269,7 @@ def test_backends_match_reference_and_agree(self): gate_cuda, up_cuda, down_cuda, + gate_up_cuda, ) torch.testing.assert_close(triton_out, ref_out, atol=0.01, rtol=0.01) @@ -282,7 +284,8 @@ def test_backends_match_reference_and_agree(self): 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) + gate_up_g = torch.cat([gate_g, up_g], dim=-1) + out_g = TritonMoeExpertsFunction.apply(num_experts, rw, selected, h_g, gate_g, up_g, down_g, gate_up_g) out_g.sum().backward() assert h_g.grad is not None and h_g.grad.abs().max() > 0 assert gate_g.grad is not None and gate_g.grad.abs().max() > 0 diff --git a/tests/optim/test_cautious_weight_decay.py b/tests/optim/test_cautious_weight_decay.py new file mode 100644 index 00000000..616de0a0 --- /dev/null +++ b/tests/optim/test_cautious_weight_decay.py @@ -0,0 +1,369 @@ +"""Tests for Cautious Weight Decay (CWD). + +Reference algorithm (Chen et al., arXiv:2510.12402):: + + x_{t+1} = x_t - eta * (u_t + lambda * I(u_t * x_t >= 0) * x_t) + +where ``u_t`` is the optimizer's update direction (post-preconditioning / +post-Newton-Schulz). When ``cautious=False`` the optimizer must reduce to its +standard decoupled-decay form. +""" + +import pytest +import torch +import torch.nn as nn +from torch.optim._muon import _adjust_lr, _zeropower_via_newtonschulz + +from xorl.optim import AnyPrecisionAdamW, Muon, SignSGD, build_optimizer +from xorl.optim.cautious import apply_cautious_decay_ + + +pytestmark = [pytest.mark.cpu] + + +# --------------------------- helper primitives ---------------------------- + + +@torch.no_grad() +def test_cautious_helper_no_op_when_weight_decay_zero(): + p = torch.tensor([1.0, -2.0]) + proxy = torch.tensor([5.0, -5.0]) + apply_cautious_decay_(p, proxy, lr=0.1, weight_decay=0.0, cautious=True) + assert torch.equal(p, torch.tensor([1.0, -2.0])) + + +@torch.no_grad() +def test_cautious_helper_matches_standard_when_cautious_false(): + p = torch.tensor([1.0, -2.0, 3.0]) + proxy = torch.tensor([1.0, -1.0, -1.0]) # mixed alignment + apply_cautious_decay_(p, proxy, lr=0.1, weight_decay=0.5, cautious=False) + expected = torch.tensor([1.0, -2.0, 3.0]) * (1 - 0.1 * 0.5) + assert torch.allclose(p, expected) + + +@torch.no_grad() +def test_cautious_helper_masks_misaligned_coordinates(): + # update * param sign: + # ( 1.0, 2.0) -> + -> decay applies + # (-1.0, 3.0) -> - -> decay skipped + # ( 0.0, -4.0) -> 0 -> decay applies (>= 0) + # ( 1.0, -5.0) -> - -> decay skipped + p = torch.tensor([2.0, 3.0, -4.0, -5.0]) + proxy = torch.tensor([1.0, -1.0, 0.0, 1.0]) + apply_cautious_decay_(p, proxy, lr=0.1, weight_decay=0.5, cautious=True) + factor = 1 - 0.1 * 0.5 + expected = torch.tensor([2.0 * factor, 3.0, -4.0 * factor, -5.0]) + assert torch.allclose(p, expected) + + +# ------------------------------ SignSGD ---------------------------------- + + +def test_signsgd_cautious_masks_decay_against_grad_sign(): + # grad * param signs: + # (+, +) -> aligned, decay applies + # (-, +) -> misaligned, decay masked + p = nn.Parameter(torch.tensor([2.0, 3.0])) + optimizer = SignSGD([p], lr=0.1, weight_decay=0.5, cautious=True) + p.grad = torch.tensor([1.0, -1.0]) + optimizer.step() + + decay = 1 - 0.1 * 0.5 + expected = torch.tensor([2.0 * decay - 0.1 * 1.0, 3.0 - 0.1 * (-1.0)]) + assert torch.allclose(p, expected) + + +def test_signsgd_cautious_equals_standard_when_signs_aligned(): + # grad and param signs both (+, -) -> all aligned -> cautious==standard. + p_std = nn.Parameter(torch.tensor([2.0, -2.0])) + p_caut = nn.Parameter(torch.tensor([2.0, -2.0])) + grad = torch.tensor([3.0, -4.0]) + + opt_std = SignSGD([p_std], lr=0.1, weight_decay=0.5, cautious=False) + opt_caut = SignSGD([p_caut], lr=0.1, weight_decay=0.5, cautious=True) + p_std.grad = grad.clone() + p_caut.grad = grad.clone() + opt_std.step() + opt_caut.step() + assert torch.allclose(p_std, p_caut) + + +# --------------------------- AnyPrecisionAdamW ---------------------------- + + +def _adamw_first_step_cautious(p_init, grad, lr, wd, cautious, beta1=0.9, beta2=0.95, eps=1e-8): + p = nn.Parameter(p_init.clone()) + optimizer = AnyPrecisionAdamW( + [p], + lr=lr, + weight_decay=wd, + betas=(beta1, beta2), + eps=eps, + momentum_dtype=torch.float32, + variance_dtype=torch.float32, + compensation_buffer_dtype=torch.float32, + cautious=cautious, + ) + p.grad = grad.clone() + optimizer.step() + return p.detach().clone() + + +def test_anyprecision_adamw_cautious_false_matches_existing_path(): + # Prior to CWD the optimizer applied decay first. Verify cautious=False + # produces the same final parameter (which is what the existing tests + # implicitly relied on). + p_init = torch.tensor([1.0, -2.0, 3.0]) + grad = torch.tensor([0.5, 0.5, -1.0]) + lr, wd = 0.1, 0.1 + beta1, beta2 = 0.9, 0.95 + eps = 1e-8 + + out = _adamw_first_step_cautious(p_init, grad, lr, wd, cautious=False, beta1=beta1, beta2=beta2, eps=eps) + + # Reference: order of operations doesn't matter for non-cautious decoupled + # decay because grad doesn't depend on p. Replicate the math. + exp_avg = (1 - beta1) * grad + exp_avg_sq = (1 - beta2) * grad * grad + bc1 = 1 - beta1 + bc2 = 1 - beta2 + denom = exp_avg_sq.sqrt() / (bc2**0.5) + eps + expected = p_init * (1 - lr * wd) - (lr / bc1) * (exp_avg / denom) + assert torch.allclose(out, expected, atol=1e-6) + + +def test_anyprecision_adamw_cautious_skips_decay_on_misaligned_coords(): + # exp_avg = (1-beta1) * grad has same sign as grad on the first step. + # Construct a case where some coords have grad sign opposite to param sign. + p_init = torch.tensor([2.0, -3.0, 4.0]) + grad = torch.tensor([-1.0, 1.0, 1.0]) # signs: -, +, + + # exp_avg signs after first step: -, +, + + # exp_avg * p signs: -, -, + + # Mask: 0, 0, 1 (decay only on coord 2) + lr, wd = 0.1, 0.5 + beta1, beta2 = 0.9, 0.95 + eps = 1e-8 + + out = _adamw_first_step_cautious(p_init, grad, lr, wd, cautious=True, beta1=beta1, beta2=beta2, eps=eps) + + exp_avg = (1 - beta1) * grad + exp_avg_sq = (1 - beta2) * grad * grad + bc1 = 1 - beta1 + bc2 = 1 - beta2 + denom = exp_avg_sq.sqrt() / (bc2**0.5) + eps + mask = (exp_avg * p_init >= 0).to(p_init.dtype) + expected = p_init * (1 - lr * wd * mask) - (lr / bc1) * (exp_avg / denom) + assert torch.allclose(out, expected, atol=1e-6) + + +def test_anyprecision_adamw_cautious_equals_standard_when_all_aligned(): + # If every coord has aligned signs, cautious must equal standard decay. + p_init = torch.tensor([1.0, 2.0, 3.0]) + grad = torch.tensor([0.5, 1.0, 0.25]) # all positive, p positive -> aligned + out_caut = _adamw_first_step_cautious(p_init, grad, 0.1, 0.5, cautious=True) + out_std = _adamw_first_step_cautious(p_init, grad, 0.1, 0.5, cautious=False) + assert torch.allclose(out_caut, out_std, atol=1e-6) + + +# -------------------------------- Muon ------------------------------------ + + +def test_muon_cautious_false_matches_standard_decay_reference(): + """``cautious=False`` must reproduce the pre-CWD Muon update exactly: + ``param *= 1 - lr*wd``, then ``param -= adjusted_lr * NS(grad)``. + + Pinned to ``standard_newton_schulz`` so the reference matches PyTorch's + upstream NS implementation; the alternate ``gram_newton_schulz`` backend + produces slightly different per-element values that would mask actual + decay-equivalence regressions. + """ + torch.manual_seed(0) + p_init = torch.randn(4, 4) + grad = torch.randn(4, 4) + lr, wd = 0.02, 0.1 + + w = nn.Parameter(p_init.clone()) + opt = Muon( + [{"params": [w], "use_muon": True}], + lr=lr, + momentum=0.0, + nesterov=False, + weight_decay=wd, + cautious=False, + ns_algorithm="standard_newton_schulz", + ) + w.grad = grad.clone() + opt.step() + + update = _zeropower_via_newtonschulz(grad, (3.4445, -4.775, 2.0315), 5, 1e-7) + adjusted_lr = _adjust_lr(lr, None, grad.shape) + expected = p_init * (1 - lr * wd) - adjusted_lr * update + assert torch.allclose(w, expected, atol=1e-4) + + +def test_muon_cautious_masks_decay_using_post_ns_update(): + """Muon's update direction is the orthogonalized matrix; the cautious + mask must be ``I(NS(grad) * param >= 0)``, not ``I(grad * param >= 0)``. + The two differ in general, so we check against an explicit reference. + """ + torch.manual_seed(0) + p_init = torch.randn(4, 4) + grad = torch.randn(4, 4) + lr, wd = 0.02, 0.5 + + # Run cautious Muon with momentum=0 so update = NS(grad). + w = nn.Parameter(p_init.clone()) + opt = Muon( + [{"params": [w], "use_muon": True}], + lr=lr, + momentum=0.0, + nesterov=False, + weight_decay=wd, + cautious=True, + ns_algorithm="standard_newton_schulz", + ) + w.grad = grad.clone() + opt.step() + actual = w.detach().clone() + + # Reference: replicate the Muon math. + update = _zeropower_via_newtonschulz(grad, (3.4445, -4.775, 2.0315), 5, 1e-7) + adjusted_lr = _adjust_lr(lr, None, grad.shape) + mask = (update * p_init >= 0).to(p_init.dtype) + expected = p_init * (1 - lr * wd * mask) - adjusted_lr * update + assert torch.allclose(actual, expected, atol=1e-4) + + +def test_muon_cautious_adamw_fallback_masks_decay(): + """The non-Muon param group (use_muon=False) takes the AdamW path; that + path must also honor cautious.""" + p_init = torch.tensor([2.0, -3.0]) + grad = torch.tensor([-1.0, 1.0]) + w = nn.Parameter(p_init.clone()) + opt = Muon( + [{"params": [w], "use_muon": False, "lr": 0.1}], + lr=0.1, + momentum=0.0, + weight_decay=0.5, + cautious=True, + ) + w.grad = grad.clone() + opt.step() + + beta1, beta2 = 0.9, 0.95 + eps = 1e-8 + exp_avg = (1 - beta1) * grad + exp_avg_sq = (1 - beta2) * grad * grad + mask = (exp_avg * p_init >= 0).to(p_init.dtype) + p_after_decay = p_init * (1 - 0.1 * 0.5 * mask) + bc1 = 1 - beta1 + bc2 = 1 - beta2 + denom = exp_avg_sq.sqrt() / (bc2**0.5) + eps + expected = p_after_decay - (0.1 / bc1) * (exp_avg / denom) + assert torch.allclose(w, expected, atol=1e-6) + + +# -------------------------- build_optimizer wiring ------------------------- + + +class _Tiny(nn.Module): + def __init__(self): + super().__init__() + self.linear = nn.Linear(3, 2) + + +def test_build_optimizer_propagates_cautious_to_signsgd(): + model = _Tiny() + opt = build_optimizer( + model, + lr=0.1, + weight_decay=0.01, + optimizer_type="signsgd", + cautious_weight_decay=True, + ) + assert isinstance(opt, SignSGD) + assert all(g["cautious"] is True for g in opt.param_groups) + + +def test_build_optimizer_propagates_cautious_to_anyprecision_adamw(): + model = _Tiny() + opt = build_optimizer( + model, + lr=0.1, + weight_decay=0.01, + optimizer_type="anyprecision_adamw", + cautious_weight_decay=True, + ) + assert isinstance(opt, AnyPrecisionAdamW) + assert all(g["cautious"] is True for g in opt.param_groups) + + +def test_build_optimizer_routes_adamw_cautious_to_anyprecision_fp32(): + model = _Tiny() + opt = build_optimizer( + model, + lr=0.1, + weight_decay=0.01, + optimizer_type="adamw", + cautious_weight_decay=True, + ) + assert isinstance(opt, AnyPrecisionAdamW) + assert all(g["cautious"] is True for g in opt.param_groups) + assert all(g["momentum_dtype"] == torch.float32 for g in opt.param_groups) + + +def test_build_optimizer_adamw_default_uses_torch_adamw_unchanged(): + model = _Tiny() + opt = build_optimizer( + model, + lr=0.1, + weight_decay=0.01, + optimizer_type="adamw", + cautious_weight_decay=False, + ) + # cautious=False keeps the fused-capable torch path + assert isinstance(opt, torch.optim.AdamW) + + +def test_build_optimizer_rejects_cautious_with_sgd(): + model = _Tiny() + with pytest.raises(ValueError, match="cautious_weight_decay is not supported"): + build_optimizer( + model, + lr=0.1, + weight_decay=0.01, + optimizer_type="sgd", + cautious_weight_decay=True, + ) + + +def test_build_optimizer_adamw_cautious_rejects_torch_adamw_only_kwargs(): + # When adamw+cautious routes to AnyPrecisionAdamW, torch.optim.AdamW-only + # kwargs (foreach, fused, amsgrad, ...) would otherwise produce a + # confusing TypeError from a class the user did not request. + model = _Tiny() + with pytest.raises(ValueError, match="routes to AnyPrecisionAdamW"): + build_optimizer( + model, + lr=0.1, + weight_decay=0.01, + optimizer_type="adamw", + cautious_weight_decay=True, + optimizer_kwargs={"foreach": False}, + ) + + +def test_build_optimizer_adamw_cautious_allows_anyprecision_kwargs(): + # use_kahan_summation is an AnyPrecisionAdamW-native kwarg and must pass + # through the cautious-routing filter. + model = _Tiny() + opt = build_optimizer( + model, + lr=0.1, + weight_decay=0.01, + optimizer_type="adamw", + cautious_weight_decay=True, + optimizer_kwargs={"use_kahan_summation": True}, + ) + assert isinstance(opt, AnyPrecisionAdamW) + assert all(g["use_kahan_summation"] is True for g in opt.param_groups) diff --git a/tests/optim/test_distsignsgd.py b/tests/optim/test_distsignsgd.py new file mode 100644 index 00000000..28fb358f --- /dev/null +++ b/tests/optim/test_distsignsgd.py @@ -0,0 +1,378 @@ +from types import SimpleNamespace + +import pytest +import torch +import torch.nn as nn + +import xorl.optim.distsignsgd as distsign_module +import xorl.optim.optimizer as optimizer_module +from xorl.optim import DistSignSGD, build_optimizer + + +pytestmark = [pytest.mark.cpu] + + +class TinyModule(nn.Module): + def __init__(self): + super().__init__() + self.linear = nn.Linear(3, 2) + + +class LocalOnlyModule(nn.Module): + def __init__(self): + super().__init__() + self.param = nn.Parameter(torch.tensor([1.0, -1.0], dtype=torch.float32)) + + +class RecordingReduceScatter: + def __init__(self): + self.seen_input = None + self.seen_op = None + self.seen_async_op = None + + def allocate(self, size, *, dtype, device): + return torch.empty(*size, dtype=dtype, device=device) + + def __call__(self, output_tensor, input_tensor, group, op, async_op=False): + self.seen_input = input_tensor.clone() + self.seen_op = op + self.seen_async_op = async_op + output_tensor.copy_(input_tensor[: output_tensor.numel()]) + return None + + +def test_distsignsgd_applies_preaggregated_gradient(): + param = nn.Parameter(torch.tensor([1.0, -2.0, 3.0])) + optimizer = DistSignSGD([param], lr=0.1) + + param.grad = torch.tensor([0.5, -0.25, 0.0]) + optimizer.step() + + expected = torch.tensor([0.95, -1.975, 3.0]) + assert torch.allclose(param, expected) + + +def test_distsignsgd_applies_decoupled_weight_decay_before_preaggregated_update(): + param = nn.Parameter(torch.tensor([2.0, -2.0])) + optimizer = DistSignSGD([param], lr=0.1, weight_decay=0.5) + + param.grad = torch.tensor([0.25, -0.5]) + optimizer.step() + + expected = torch.tensor([1.875, -1.85]) + assert torch.allclose(param, expected) + + +def test_distsignsgd_rejects_sparse_gradients(): + param = nn.Parameter(torch.ones(4)) + optimizer = DistSignSGD([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_distsignsgd_keeps_optimizer_state_empty_across_steps(): + param = nn.Parameter(torch.tensor([1.0])) + optimizer = DistSignSGD([param], lr=0.1) + + for grad in (torch.tensor([1.0]), torch.tensor([-0.5])): + param.grad = grad + optimizer.step() + assert len(optimizer.state) == 0 + + +def test_distsignsgd_state_dict_round_trips_without_state_tensors(): + source_param = nn.Parameter(torch.tensor([1.0, -1.0])) + source_optimizer = DistSignSGD([source_param], lr=0.1, weight_decay=0.25) + source_param.grad = torch.tensor([0.5, -0.25]) + source_optimizer.step() + + state_dict = source_optimizer.state_dict() + + target_param = nn.Parameter(torch.tensor([0.0, 0.0])) + target_optimizer = DistSignSGD([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_dist_sign_reduce_scatter_signs_input_before_reduce(): + inner = RecordingReduceScatter() + comm = distsign_module.DistSignReduceScatter(inner_comm=inner) + + input_tensor = torch.tensor([2.0, -0.5, 0.0], dtype=torch.float32) + output_tensor = torch.empty_like(input_tensor) + + comm( + output_tensor=output_tensor, + input_tensor=input_tensor, + group=None, + op=torch.distributed.ReduceOp.SUM, + async_op=True, + ) + + expected = torch.tensor([1.0, -1.0, 0.0], dtype=torch.float32) + assert torch.equal(input_tensor, expected) + assert torch.equal(inner.seen_input, expected) + assert torch.equal(output_tensor, expected) + assert inner.seen_op == torch.distributed.ReduceOp.SUM + assert inner.seen_async_op is True + + +def test_dist_sign_reduce_scatter_forces_sum_when_caller_passes_avg(): + inner = RecordingReduceScatter() + comm = distsign_module.DistSignReduceScatter(inner_comm=inner) + + input_tensor = torch.tensor([2.0, -0.5, 0.0], dtype=torch.float32) + output_tensor = torch.empty_like(input_tensor) + + comm( + output_tensor=output_tensor, + input_tensor=input_tensor, + group=None, + op=torch.distributed.ReduceOp.AVG, + ) + + # FSDP2 with reduce_dtype=fp32 may pass AVG; the trainer's voter-total + # divisor would then double-divide. Forcing SUM here keeps the sign-vote + # accumulator semantics intact. + assert inner.seen_op == torch.distributed.ReduceOp.SUM + + +def test_dist_sign_reduce_scatter_sums_sp_before_sign(monkeypatch): + inner = RecordingReduceScatter() + comm = distsign_module.DistSignReduceScatter(inner_comm=inner, sp_group="sp-group") + reduced = [] + + def fake_all_reduce(tensor, op, group): + reduced.append((tensor.clone(), op, group)) + tensor.add_(torch.tensor([-3.0, 1.0, 0.0], dtype=tensor.dtype)) + + monkeypatch.setattr(distsign_module.dist, "all_reduce", fake_all_reduce) + + input_tensor = torch.tensor([2.0, -0.5, 0.0], dtype=torch.float32) + output_tensor = torch.empty_like(input_tensor) + + comm( + output_tensor=output_tensor, + input_tensor=input_tensor, + group=None, + op=torch.distributed.ReduceOp.SUM, + ) + + expected = torch.tensor([-1.0, 1.0, 0.0], dtype=torch.float32) + assert len(reduced) == 1 + assert reduced[0][1] == torch.distributed.ReduceOp.SUM + assert reduced[0][2] == "sp-group" + assert torch.equal(input_tensor, expected) + assert torch.equal(inner.seen_input, expected) + assert torch.equal(output_tensor, expected) + + +def test_configure_distsignsgd_registers_local_sign_hook(monkeypatch): + model = LocalOnlyModule() + + monkeypatch.setattr( + distsign_module, + "get_parallel_state", + lambda: SimpleNamespace( + dp_mode="fsdp2", + dp_replicate_enabled=False, + ep_enabled=False, + cp_enabled=False, + cp_fsdp_mode="none", + sp_grad_sync_group=None, + ), + ) + + distsign_module.configure_distsignsgd(model) + + (model.param * torch.tensor([2.0, -3.0])).sum().backward() + (model.param * torch.tensor([-4.0, 5.0])).sum().backward() + + assert torch.equal(model.param.grad, torch.tensor([0.0, 0.0])) + assert getattr(model.param, "_distsign_local_hook_registered", False) is True + assert getattr(model, "_distsignsgd_configured", False) is True + + +def test_configure_distsignsgd_skips_local_sign_hook_for_fsdp_managed_params(monkeypatch): + class FakeFSDPModule(nn.Module): + def __init__(self): + super().__init__() + self.managed = nn.Parameter(torch.tensor([1.0, -1.0], dtype=torch.float32)) + + def _get_fsdp_state(self): + return SimpleNamespace( + _fsdp_param_group=SimpleNamespace( + fsdp_params=[SimpleNamespace(sharded_param=self.managed)], + ) + ) + + def set_custom_reduce_scatter(self, comm): + self._reduce_scatter_comm = comm + + class MixedModel(nn.Module): + def __init__(self): + super().__init__() + self.fsdp = FakeFSDPModule() + self.local = nn.Parameter(torch.tensor([1.0, -1.0], dtype=torch.float32)) + + model = MixedModel() + + monkeypatch.setattr(distsign_module, "FSDPModule", FakeFSDPModule) + monkeypatch.setattr( + distsign_module, + "get_parallel_state", + lambda: SimpleNamespace( + dp_mode="fsdp2", + dp_replicate_enabled=False, + ep_enabled=False, + cp_enabled=False, + cp_fsdp_mode="none", + sp_grad_sync_group=None, + ), + ) + + distsign_module.configure_distsignsgd(model) + + (model.fsdp.managed * torch.tensor([2.0, -3.0])).sum().backward() + (model.fsdp.managed * torch.tensor([-4.0, 5.0])).sum().backward() + (model.local * torch.tensor([2.0, -3.0])).sum().backward() + (model.local * torch.tensor([-4.0, 5.0])).sum().backward() + + assert torch.equal(model.fsdp.managed.grad, torch.tensor([-2.0, 2.0])) + assert torch.equal(model.local.grad, torch.tensor([0.0, 0.0])) + assert getattr(model.local, "_distsign_local_hook_registered", False) is True + assert getattr(model.fsdp.managed, "_distsign_local_hook_registered", False) is False + + +def test_configure_distsignsgd_rejects_hsdp(monkeypatch): + model = LocalOnlyModule() + + monkeypatch.setattr( + distsign_module, + "get_parallel_state", + lambda: SimpleNamespace(dp_mode="fsdp2", dp_replicate_enabled=True), + ) + + with pytest.raises(NotImplementedError, match="does not yet support HSDP"): + distsign_module.configure_distsignsgd(model) + + +def test_configure_distsignsgd_rejects_sequence_parallel_folded_into_fsdp(monkeypatch): + model = LocalOnlyModule() + + monkeypatch.setattr( + distsign_module, + "get_parallel_state", + lambda: SimpleNamespace( + dp_mode="fsdp2", + dp_replicate_enabled=False, + ep_enabled=False, + cp_enabled=True, + cp_fsdp_mode="all", + ), + ) + + with pytest.raises(NotImplementedError, match="set cp_fsdp_mode='none'"): + distsign_module.configure_distsignsgd(model) + + +def test_configure_distsignsgd_rejects_expert_parallelism(monkeypatch): + model = LocalOnlyModule() + + monkeypatch.setattr( + distsign_module, + "get_parallel_state", + lambda: SimpleNamespace( + dp_mode="fsdp2", + dp_replicate_enabled=False, + ep_enabled=True, + cp_enabled=False, + cp_fsdp_mode="none", + ), + ) + + with pytest.raises(NotImplementedError, match="expert parallelism"): + distsign_module.configure_distsignsgd(model) + + +def test_configure_distsignsgd_rejects_non_fsdp_dtensor_param(monkeypatch): + class FakeDTensor: + pass + + class TPOnlyModule(nn.Module): + def __init__(self): + super().__init__() + tp_param = nn.Parameter(torch.tensor([1.0, -1.0], dtype=torch.float32)) + # Spoof a non-FSDP DTensor parameter. + tp_param.__class__ = type("DTensorParam", (FakeDTensor, nn.Parameter), {}) + self.tp_param = tp_param + + model = TPOnlyModule() + + monkeypatch.setattr(distsign_module, "DTensor", FakeDTensor) + monkeypatch.setattr( + distsign_module, + "get_parallel_state", + lambda: SimpleNamespace( + dp_mode="fsdp2", + dp_replicate_enabled=False, + ep_enabled=False, + cp_enabled=False, + cp_fsdp_mode="none", + sp_grad_sync_group=None, + ), + ) + + with pytest.raises(NotImplementedError, match="DTensor parameter that is not FSDP-managed"): + distsign_module.configure_distsignsgd(model) + + +def test_build_optimizer_supports_distsignsgd_and_preserves_weight_decay_split(monkeypatch): + model = TinyModule() + configured = [] + + monkeypatch.setattr( + optimizer_module, + "get_parallel_state", + lambda: SimpleNamespace(dp_mode="fsdp2", ep_enabled=False), + ) + monkeypatch.setattr( + optimizer_module, "configure_distsignsgd", lambda configured_model: configured.append(configured_model) + ) + + optimizer = build_optimizer( + model, + lr=0.1, + weight_decay=0.01, + optimizer_type="distsignsgd", + no_decay_params=["bias"], + ) + + assert isinstance(optimizer, DistSignSGD) + assert configured == [model] + 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] + + +def test_build_optimizer_rejects_distsignsgd_without_fsdp2(monkeypatch): + model = TinyModule() + + monkeypatch.setattr( + optimizer_module, + "get_parallel_state", + lambda: SimpleNamespace(dp_mode="ddp"), + ) + + with pytest.raises(ValueError, match="requires data_parallel_mode='fsdp2'"): + build_optimizer(model, optimizer_type="distsignsgd") diff --git a/tests/optim/test_lr_scheduler.py b/tests/optim/test_lr_scheduler.py new file mode 100644 index 00000000..635de5e9 --- /dev/null +++ b/tests/optim/test_lr_scheduler.py @@ -0,0 +1,153 @@ +import pytest +import torch.nn as nn +from torch.optim import SGD + +from xorl.optim.lr_scheduler import build_lr_scheduler + + +pytestmark = [pytest.mark.cpu] + + +def _trace(scheduler, steps: int) -> list[float]: + lrs: list[float] = [] + for _ in range(steps): + lrs.append(scheduler.get_last_lr()[0]) + scheduler.step() + return lrs + + +def _make_optimizer(lr: float = 1.0) -> SGD: + return SGD(nn.Linear(2, 2).parameters(), lr=lr) + + +class TestConstantSchedule: + def test_no_warmup_holds_lr(self): + sched = build_lr_scheduler(_make_optimizer(), train_steps=8, lr=1.0, lr_decay_style="constant") + assert _trace(sched, 8) == pytest.approx([1.0] * 8) + + def test_linear_warmup_then_constant(self): + sched = build_lr_scheduler( + _make_optimizer(), + train_steps=10, + lr=1.0, + lr_decay_style="constant", + lr_warmup_ratio=0.4, + lr_start=0.0, + ) + lrs = _trace(sched, 8) + # 4 warmup steps from lr_start=0 to init_lr=1, then constant at 1 + assert lrs[:4] == pytest.approx([0.0, 0.25, 0.5, 0.75]) + assert lrs[4:] == pytest.approx([1.0] * 4) + + +class TestLinearSchedule: + def test_default_decays_to_near_zero_over_full_range(self): + sched = build_lr_scheduler(_make_optimizer(), train_steps=10, lr=1.0, lr_decay_style="linear") + lrs = _trace(sched, 11) + # No warmup, default lr_min=1e-7, decay_ratio=1: linear 1.0 → ~0 over 10 steps. + # Last value at step 10 is min_lr_ratio = 1e-7. + assert lrs[0] == pytest.approx(1.0) + assert lrs[-1] == pytest.approx(1e-7) + diffs = [lrs[i] - lrs[i + 1] for i in range(len(lrs) - 2)] + assert all(d == pytest.approx(diffs[0], rel=1e-6) for d in diffs) + + def test_lr_min_floor(self): + sched = build_lr_scheduler( + _make_optimizer(), + train_steps=10, + lr=1.0, + lr_decay_style="linear", + lr_min=0.25, + ) + lrs = _trace(sched, 12) + assert lrs[0] == pytest.approx(1.0) + assert lrs[-1] == pytest.approx(0.25) + assert min(lrs) == pytest.approx(0.25) + + def test_warmup_then_decay_to_lr_min_within_decay_ratio(self): + sched = build_lr_scheduler( + _make_optimizer(), + train_steps=10, + lr=1.0, + lr_decay_style="linear", + lr_warmup_ratio=0.2, + lr_min=0.1, + lr_decay_ratio=0.8, + ) + lrs = _trace(sched, 12) + # warmup steps 0-1 (lr_start=0 to 1), decay steps 2-7 (1.0 → 0.1), floor 8+. + assert lrs[0] == pytest.approx(0.0) + assert lrs[1] == pytest.approx(0.5) + assert lrs[2] == pytest.approx(1.0) + assert lrs[7] == pytest.approx(0.25) + assert lrs[8] == pytest.approx(0.1) + assert lrs[11] == pytest.approx(0.1) + + +class TestCosineSchedule: + def test_endpoints_and_midpoint(self): + sched = build_lr_scheduler( + _make_optimizer(), + train_steps=10, + lr=1.0, + lr_decay_style="cosine", + lr_min=0.0, + ) + lrs = _trace(sched, 11) + assert lrs[0] == pytest.approx(1.0) + # half-cosine: at progress=0.5, factor = 0.5 + assert lrs[5] == pytest.approx(0.5, abs=1e-6) + assert lrs[10] == pytest.approx(0.0, abs=1e-6) + + def test_lr_min_floor_after_decay_ratio(self): + sched = build_lr_scheduler( + _make_optimizer(), + train_steps=10, + lr=1.0, + lr_decay_style="cosine", + lr_min=0.1, + lr_decay_ratio=0.8, + ) + lrs = _trace(sched, 12) + assert lrs[0] == pytest.approx(1.0) + assert lrs[8] == pytest.approx(0.1) + assert lrs[11] == pytest.approx(0.1) + # Monotonically non-increasing within decay window + for a, b in zip(lrs[:9], lrs[1:9]): + assert b <= a + 1e-9 + + def test_warmup_then_cosine(self): + sched = build_lr_scheduler( + _make_optimizer(), + train_steps=10, + lr=1.0, + lr_decay_style="cosine", + lr_warmup_ratio=0.2, + lr_min=0.0, + ) + lrs = _trace(sched, 11) + # 2 warmup steps from 0 → 1, then half-cosine from 1 → 0 over 8 steps. + assert lrs[0] == pytest.approx(0.0) + assert lrs[1] == pytest.approx(0.5) + assert lrs[2] == pytest.approx(1.0) + # midpoint of decay (step 6): cos(π/2) → factor 0.5 + assert lrs[6] == pytest.approx(0.5, abs=1e-6) + assert lrs[10] == pytest.approx(0.0, abs=1e-6) + + +class TestValidation: + def test_rejects_non_positive_lr(self): + with pytest.raises(ValueError, match="lr must be > 0"): + build_lr_scheduler(_make_optimizer(), train_steps=10, lr=0.0) + with pytest.raises(ValueError, match="lr must be > 0"): + build_lr_scheduler(_make_optimizer(), train_steps=10, lr=-1e-3) + + def test_rejects_warmup_ratio_out_of_range(self): + with pytest.raises(ValueError, match="lr_warmup_ratio"): + build_lr_scheduler(_make_optimizer(), train_steps=10, lr=1.0, lr_warmup_ratio=1.5) + with pytest.raises(ValueError, match="lr_warmup_ratio"): + build_lr_scheduler(_make_optimizer(), train_steps=10, lr=1.0, lr_warmup_ratio=-0.1) + + def test_rejects_unknown_decay_style(self): + with pytest.raises(ValueError, match="Unknown learning rate decay style"): + build_lr_scheduler(_make_optimizer(), train_steps=10, lr=1.0, lr_decay_style="bogus") diff --git a/tests/optim/test_muon.py b/tests/optim/test_muon.py index 0b3ff118..9a5576b5 100644 --- a/tests/optim/test_muon.py +++ b/tests/optim/test_muon.py @@ -76,6 +76,7 @@ def test_build_optimizer_threads_gram_newton_schulz_kwargs(): "muon_ns_algorithm": "gram_newton_schulz", "muon_ns_use_quack_kernels": False, "muon_gram_ns_num_restarts": 1, + "muon_grouped_gram_ns_fp32_byte_limit": 23, "muon_grad_dtype": "fp32", "muon_update_dtype": "bf16", "muon_force_momentum_path": True, @@ -91,12 +92,20 @@ def test_build_optimizer_threads_gram_newton_schulz_kwargs(): 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 all(group["grouped_gram_ns_fp32_byte_limit"] == 23 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_muon_rejects_nonpositive_grouped_gram_newton_schulz_byte_limit(): + param = nn.Parameter(torch.zeros((2, 3), dtype=torch.float32)) + + with pytest.raises(ValueError, match="grouped_gram_ns_fp32_byte_limit must be positive"): + Muon([param], grouped_gram_ns_fp32_byte_limit=0) + + def test_make_quack_backend_prefers_installed_quack(monkeypatch): fake_gemm_interface = SimpleNamespace( gemm=lambda A, B: ("gemm", A, B), @@ -344,13 +353,13 @@ def orthogonalize(self, X): 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 + def fake_batched_zeropower(update, ns_coefficients, ns_steps, eps): + # Receives the flattened-batch tensor [B, H, I]; assign each batch element + # a distinct constant offset so we can confirm correct un-flattening. seen_shapes.append(tuple(update.shape)) - return update + call_count + offsets = torch.arange(1, update.shape[0] + 1, dtype=update.dtype, device=update.device).reshape(-1, 1, 1) + return update + offsets p = nn.Parameter(torch.zeros((2, 3, 4), dtype=torch.float32)) optimizer = Muon( @@ -362,13 +371,14 @@ def fake_zeropower(update, ns_coefficients, ns_steps, eps): ) monkeypatch.setattr(muon_module, "_adjust_lr", lambda lr, adjust_lr_fn, shape: lr) - monkeypatch.setattr(muon_module, "_zeropower_via_newtonschulz", fake_zeropower) + monkeypatch.setattr(muon_module, "_batched_zeropower_via_newtonschulz", fake_batched_zeropower) p.grad = torch.ones((2, 3, 4), dtype=torch.float32) optimizer.step() - assert seen_shapes == [(3, 4), (3, 4)] + # Single batched call over the flattened leading dims: [B=2, H=3, I=4]. + assert seen_shapes == [(2, 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]], @@ -397,11 +407,11 @@ def orthogonalize(self, X): nesterov=False, ns_algorithm="gram_newton_schulz", ns_use_quack_kernels=False, + grouped_gram_ns_fp32_byte_limit=23, ) 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) diff --git a/tests/optim/test_stochastic_round.py b/tests/optim/test_stochastic_round.py new file mode 100644 index 00000000..a1a4ca56 --- /dev/null +++ b/tests/optim/test_stochastic_round.py @@ -0,0 +1,72 @@ +"""Tests for stochastic rounding to BF16.""" + +import pytest +import torch + +from xorl.optim.stochastic_round import stochastic_round_to_bf16 + + +pytestmark = [pytest.mark.cpu] + + +def test_stochastic_round_dtype_and_shape(): + x = torch.randn(7, 13, dtype=torch.float32) + y = stochastic_round_to_bf16(x) + assert y.dtype == torch.bfloat16 + assert y.shape == x.shape + assert y.device == x.device + + +def test_stochastic_round_rejects_non_fp32(): + x = torch.randn(4, 4, dtype=torch.bfloat16) + with pytest.raises(ValueError): + stochastic_round_to_bf16(x) + + +def test_stochastic_round_is_unbiased_in_expectation(): + """E[round(x)] should approach x as we average many samples.""" + torch.manual_seed(0) + x = torch.randn(64, 64, dtype=torch.float32) + n_samples = 4000 + accum = torch.zeros_like(x) + for _ in range(n_samples): + accum += stochastic_round_to_bf16(x).to(torch.float32) + mean = accum / n_samples + # With n=4000 samples and BF16 ulp ~ |x| * 2**-7, the standard error of + # the mean is ~ulp / sqrt(n) ≈ ulp / 63. The test allows a generous + # multiple to keep it deterministic across platforms. + rel_err = ((mean - x).abs() / (x.abs() + 1e-8)).max().item() + assert rel_err < 1e-2, f"Mean relative error {rel_err} exceeds tolerance; expected unbiased" + + +def test_stochastic_round_within_neighbors(): + """Output is always one of the two BF16 neighbors of x (no overshoot).""" + torch.manual_seed(1) + x = torch.randn(256, dtype=torch.float32) + # The two BF16 neighbors of an FP32 x are obtained by truncate-down (mask + # off low bits) and truncate-down + 1 ulp. + x_int = x.view(torch.int32) + lower_int = x_int & ~0xFFFF + upper_int = lower_int + 0x10000 + lower = lower_int.view(torch.float32).to(torch.bfloat16).to(torch.float32) + upper = upper_int.view(torch.float32).to(torch.bfloat16).to(torch.float32) + # Note: when x is exactly representable in BF16, lower == x; upper is the + # next BF16 above, but stochastic round only ever returns x in that case. + for _ in range(20): + y = stochastic_round_to_bf16(x).to(torch.float32) + # y must match either lower or upper for every element + match_lower = y == lower + match_upper = y == upper + match_x = y == x # exact-representation case + assert (match_lower | match_upper | match_x).all(), ( + "stochastic rounded value is not one of the two BF16 neighbors" + ) + + +def test_stochastic_round_deterministic_with_generator(): + x = torch.randn(32, dtype=torch.float32) + g1 = torch.Generator().manual_seed(42) + g2 = torch.Generator().manual_seed(42) + y1 = stochastic_round_to_bf16(x, generator=g1) + y2 = stochastic_round_to_bf16(x, generator=g2) + assert torch.equal(y1, y2) diff --git a/tests/qlora/test_quantize_error_reduction.py b/tests/qlora/test_quantize_error_reduction.py deleted file mode 100644 index d1f9ced6..00000000 --- a/tests/qlora/test_quantize_error_reduction.py +++ /dev/null @@ -1,6 +0,0 @@ -"""Tests for QLoRA quantization error reduction techniques.""" - -import pytest - - -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 a3d46897..719567f4 100644 --- a/tests/server/api_server/test_api_server.py +++ b/tests/server/api_server/test_api_server.py @@ -7,28 +7,83 @@ import asyncio import time +import types +from types import SimpleNamespace import pytest +from fastapi import HTTPException from xorl.server.api_server.api_types import ( AdamParams, + CreateModelRequest, CreateSessionRequest, Datum, DatumInput, ForwardBackwardRequest, LoadWeightsRequest, OptimStepRequest, + SaveWeightsForSamplerRequest, SaveWeightsRequest, SessionHeartbeatRequest, UntypedAPIFuture, + WeightsInfoRequest, +) +from xorl.server.api_server.endpoints import ( + create_model_endpoint, + create_session_endpoint, + save_weights_endpoint, + session_heartbeat_endpoint, + weights_info_endpoint, ) -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 +from xorl.server.protocol.operations import RegisterSessionData +from xorl.server.session_spec import build_default_session_spec, write_session_spec pytestmark = [pytest.mark.cpu, pytest.mark.server] +class _ImmediateFutureStore: + def __init__(self) -> None: + self.last_result = None + + async def create(self, *, model_id, request_type, process_fn, request_data, ttl=None): + self.last_result = await process_fn(request_data) + return "future-test-1" + + +class _FakeOrchestratorClient: + def __init__(self) -> None: + self.last_request = None + + async def send_request(self, request): + self.last_request = request + return request + + +def _build_default_session_spec(): + train_config = { + "optimizer": "adamw", + "lr": 1e-5, + "weight_decay": 0.01, + "optimizer_dtype": "bf16", + "optimizer_kwargs": {}, + } + lora_config = { + "enable_lora": True, + "lora_rank": 8, + "max_lora_rank": 16, + "lora_alpha": 16, + "lora_target_modules": ["q_proj", "o_proj"], + } + default_session_spec = build_default_session_spec( + base_model="Qwen/Qwen3-8B", + train_config=train_config, + lora_config=lora_config, + ) + return default_session_spec, lora_config + + class TestAPIServerConfiguration: """Test APIServer configuration and initialization.""" @@ -93,13 +148,17 @@ def test_request_creation_and_serialization(self): # OptimStepRequest request = OptimStepRequest(adam_params=AdamParams(learning_rate=1e-4), gradient_clip=1.0) - assert request.adam_params.learning_rate == 1e-4 + assert request.learning_rate == 1e-4 + assert request.adam_params is not None and request.adam_params.learning_rate == 1e-4 data = request.model_dump() + assert data["learning_rate"] == 1e-4 assert data["adam_params"]["learning_rate"] == 1e-4 # Defaults request = OptimStepRequest() - assert request.adam_params.learning_rate == 0.0001 and request.gradient_clip is None + assert request.learning_rate is None + assert request.adam_params is None + assert request.gradient_clip is None # SaveWeightsRequest request = SaveWeightsRequest(path="/tmp/checkpoint") @@ -111,6 +170,51 @@ def test_request_creation_and_serialization(self): request = LoadWeightsRequest(path="/tmp/checkpoint", optimizer=True) assert request.optimizer is True + def test_optim_step_learning_rate_resolution(self): + server = APIServer( + engine_input_addr="tcp://127.0.0.1:17004", + engine_output_addr="tcp://127.0.0.1:17005", + ) + server.model_configs["full-session"] = {"base_model": "base", "lora_config": {}} + server.model_configs["lora-session"] = { + "base_model": "base", + "lora_config": {"lora_rank": 8}, + "optimizer_config": {"learning_rate": 3e-5}, + } + + assert server._optim_step_learning_rate(OptimStepRequest(model_id="full-session")) == AdamParams().learning_rate + assert server._optim_step_learning_rate(OptimStepRequest(model_id="lora-session")) == 3e-5 + assert ( + server._optim_step_learning_rate( + OptimStepRequest(model_id="lora-session", adam_params=AdamParams(learning_rate=2e-4)) + ) + == 2e-4 + ) + assert server._optim_step_learning_rate(OptimStepRequest(model_id="lora-session", learning_rate=7e-5)) == 7e-5 + + class FakeClient: + request = None + + async def send_request(self, request): + self.request = request + return object() + + async def fake_wait_for_response(response_future, request_id, timeout, message): + class Output: + outputs = [{"grad_norm": 0.0}] + + return Output() + + fake_client = FakeClient() + server._running = True + server.orchestrator_client = fake_client + server._wait_for_response = fake_wait_for_response + + response = asyncio.run(server.optim_step(OptimStepRequest(model_id="full-session"))) + + assert fake_client.request.payload.lr == AdamParams().learning_rate + assert response.metrics["learning_rate"] == AdamParams().learning_rate + class TestTinkerSessionCompatibility: """Test Tinker-compatible session creation and heartbeats.""" @@ -178,6 +282,424 @@ def test_session_heartbeat_refreshes_activity(self): assert heartbeat_response.session_id == session_id assert server.session_last_activity[session_id] > initial_activity + def test_create_session_stores_canonical_lora_config(self): + """Tinker rank/alpha aliases should be canonicalized before server storage.""" + server = APIServer( + engine_input_addr="tcp://127.0.0.1:17022", + engine_output_addr="tcp://127.0.0.1:17023", + ) + + response = asyncio.run( + create_session_endpoint( + CreateSessionRequest(session_id="alias-session", lora_config={"rank": 6, "alpha": 14}), + server=server, + ) + ) + + assert response.session_id == "alias-session" + assert server.model_configs["alias-session"]["lora_config"] == { + "lora_rank": 6, + "lora_alpha": 14, + } + + +class TestTinkerCompatibilityPaths: + """Exercise HTTP-boundary compatibility paths, not just Pydantic parsing.""" + + def test_weights_info_accepts_and_serializes_legacy_lora_rank(self, tmp_path): + """weights_info should keep Tinker's flat lora_rank response field.""" + default_session_spec, _ = _build_default_session_spec() + server = APIServer( + engine_input_addr="tcp://127.0.0.1:17014", + engine_output_addr="tcp://127.0.0.1:17015", + output_dir=str(tmp_path), + base_model="Qwen/Qwen3-8B", + ) + checkpoint_path = tmp_path / "weights" / "session-a" / "checkpoint-001" + write_session_spec(str(checkpoint_path), default_session_spec) + + response = asyncio.run( + weights_info_endpoint( + WeightsInfoRequest(xorl_path="xorl://session-a/weights/checkpoint-001"), + server=server, + ) + ) + + assert response.base_model == "Qwen/Qwen3-8B" + assert response.is_lora is True + assert response.lora_rank == 8 + assert "lora_rank" in response.model_dump() + + def test_create_model_stores_dict_lora_config_for_weights_info(self, tmp_path): + """create_model should not store a typed LoRAConfigRequest where weights_info expects a dict.""" + default_session_spec, lora_config = _build_default_session_spec() + server = APIServer( + engine_input_addr="tcp://127.0.0.1:17016", + engine_output_addr="tcp://127.0.0.1:17017", + output_dir=str(tmp_path), + base_model="Qwen/Qwen3-8B", + default_session_spec=default_session_spec, + server_lora_config=lora_config, + max_lora_rank=16, + skip_initial_checkpoint=True, + ) + server.future_store = _ImmediateFutureStore() + server.orchestrator_client = _FakeOrchestratorClient() + server._running = True + + async def _wait_for_response(self, response_future, request_id, timeout, timeout_message="timeout"): + return SimpleNamespace( + outputs={"result": {"registered": True, "model_id": response_future.payload.model_id}} + ) + + server._wait_for_response = types.MethodType(_wait_for_response, server) + + create_response = asyncio.run( + create_model_endpoint( + CreateModelRequest( + model_id="session-a", + base_model="Qwen/Qwen3-8B", + lora_config={"rank": 8, "alpha": 16}, + optimizer_config={"type": "adamw", "learning_rate": 4e-5}, + ), + server=server, + ) + ) + + assert create_response.request_id == "future-test-1" + assert server.future_store.last_result == {"model_id": "session-a", "type": "create_model"} + assert server.model_configs["session-a"]["lora_config"] == {"lora_rank": 8, "lora_alpha": 16} + assert server.model_configs["session-a"]["optimizer_config"]["type"] == "adamw" + assert server.model_configs["session-a"]["optimizer_config"]["learning_rate"] == 4e-5 + checkpoint_path = tmp_path / "weights" / "session-a" / "checkpoint-001" + write_session_spec(str(checkpoint_path), server.model_configs["session-a"]) + + info_response = asyncio.run( + weights_info_endpoint( + WeightsInfoRequest(xorl_path="xorl://session-a/weights/checkpoint-001"), + server=server, + ) + ) + + assert info_response.base_model == "Qwen/Qwen3-8B" + assert info_response.lora_rank == 8 + assert info_response.model_dump()["lora_rank"] == 8 + + def test_full_weight_create_model_allows_empty_optional_configs(self): + """Empty lora/optimizer configs from legacy clients are no-op full-weight overrides.""" + server = APIServer( + engine_input_addr="tcp://127.0.0.1:17032", + engine_output_addr="tcp://127.0.0.1:17033", + base_model="Qwen/Qwen3-8B", + skip_initial_checkpoint=True, + train_config={"lr": 5e-5}, + lora_config={"enable_lora": False}, + ) + server.future_store = _ImmediateFutureStore() + client = _FakeOrchestratorClient() + server._running = True + server.orchestrator_client = client + + async def _wait_for_response(self, response_future, request_id, timeout, timeout_message="Engine timeout"): + assert timeout_message == "Register session timeout" + return SimpleNamespace(error=None, outputs=[{"result": {"registered": True}}]) + + server._wait_for_response = types.MethodType(_wait_for_response, server) + + create_response = asyncio.run( + create_model_endpoint( + CreateModelRequest( + model_id="default", + base_model="Qwen/Qwen3-8B", + lora_config={}, + optimizer_config={}, + ), + server=server, + ) + ) + + assert create_response.request_id == "future-test-1" + assert client.last_request.payload.session_spec == { + "base_model": "Qwen/Qwen3-8B", + "is_lora": False, + } + assert client.last_request.payload.materialize is False + + def test_create_model_registers_lora_session_on_workers(self): + """LoRA create_model must register a worker session spec before exposing the model_id.""" + train_config = { + "optimizer": "adamw", + "lr": 2e-5, + "weight_decay": 0.02, + "optimizer_dtype": "bf16", + "optimizer_kwargs": {}, + } + lora_config = { + "enable_lora": True, + "lora_rank": 4, + "max_lora_rank": 16, + "lora_alpha": 8, + } + default_session_spec = build_default_session_spec( + base_model="Qwen/Qwen3-8B", + train_config=train_config, + lora_config=lora_config, + ) + server = APIServer( + engine_input_addr="tcp://127.0.0.1:17018", + engine_output_addr="tcp://127.0.0.1:17019", + base_model="Qwen/Qwen3-8B", + skip_initial_checkpoint=True, + default_session_spec=default_session_spec, + server_lora_config=lora_config, + max_lora_rank=16, + train_config=train_config, + lora_config=lora_config, + ) + server.future_store = _ImmediateFutureStore() + client = _FakeOrchestratorClient() + server._running = True + server.orchestrator_client = client + + async def _wait_for_response(self, response_future, request_id, timeout, timeout_message="Engine timeout"): + assert response_future is client.last_request + assert request_id == client.last_request.request_id + assert timeout_message == "Register session timeout" + return SimpleNamespace(error=None, outputs=[{"result": {"registered": True}}]) + + server._wait_for_response = types.MethodType(_wait_for_response, server) + + create_response = asyncio.run( + create_model_endpoint( + CreateModelRequest( + model_id="policy-a", + base_model="Qwen/Qwen3-8B", + lora_config={"rank": 12, "alpha": 24}, + optimizer_config={"learning_rate": 3e-5}, + ), + server=server, + ) + ) + + assert create_response.request_id == "future-test-1" + assert client.last_request.operation == "register_session" + assert isinstance(client.last_request.payload, RegisterSessionData) + payload = client.last_request.payload + assert payload.model_id == "policy-a" + assert payload.materialize is True + assert payload.session_spec["lora_config"] == {"lora_rank": 12, "lora_alpha": 24} + assert payload.session_spec["optimizer_config"]["learning_rate"] == pytest.approx(3e-5) + assert payload.session_spec["optimizer_config"]["weight_decay"] == pytest.approx(0.02) + assert server.model_configs["policy-a"]["lora_config"] == {"lora_rank": 12, "lora_alpha": 24} + assert server.model_configs["policy-a"]["optimizer_config"]["learning_rate"] == pytest.approx(3e-5) + + def test_optim_step_supports_native_payload_without_adam_params(self): + """Native xorl clients can send learning_rate/gradient_clip without legacy adam_params.""" + server = APIServer( + engine_input_addr="tcp://127.0.0.1:17018", + engine_output_addr="tcp://127.0.0.1:17019", + ) + client = _FakeOrchestratorClient() + server._running = True + server.orchestrator_client = client + + async def _wait_for_response(self, response_future, request_id, timeout, timeout_message="timeout"): + return SimpleNamespace(outputs=[{"grad_norm": 0.5}]) + + server._wait_for_response = types.MethodType(_wait_for_response, server) + + response = asyncio.run( + server.optim_step( + OptimStepRequest( + model_id="default", + learning_rate=3e-4, + gradient_clip=1.25, + ) + ) + ) + + assert client.last_request.operation == "optim_step" + assert client.last_request.payload.lr == pytest.approx(3e-4) + assert client.last_request.payload.gradient_clip == pytest.approx(1.25) + assert client.last_request.payload.beta1 is None + assert response.metrics["learning_rate"] == pytest.approx(3e-4) + + def test_optim_step_supports_tinker_adam_params_payload(self): + """Legacy Tinker adam_params should still drive lr, clip, and Adam hyperparameters.""" + server = APIServer( + engine_input_addr="tcp://127.0.0.1:17020", + engine_output_addr="tcp://127.0.0.1:17021", + ) + client = _FakeOrchestratorClient() + server._running = True + server.orchestrator_client = client + + async def _wait_for_response(self, response_future, request_id, timeout, timeout_message="timeout"): + return SimpleNamespace(outputs=[{"grad_norm": 0.5}]) + + server._wait_for_response = types.MethodType(_wait_for_response, server) + + response = asyncio.run( + server.optim_step( + OptimStepRequest( + model_id="default", + adam_params=AdamParams( + learning_rate=2e-4, + beta1=0.8, + beta2=0.88, + eps=1e-6, + grad_clip_norm=0.75, + ), + ) + ) + ) + + assert client.last_request.operation == "optim_step" + assert client.last_request.payload.lr == pytest.approx(2e-4) + assert client.last_request.payload.gradient_clip == pytest.approx(0.75) + assert client.last_request.payload.beta1 == pytest.approx(0.8) + assert client.last_request.payload.beta2 == pytest.approx(0.88) + assert client.last_request.payload.eps == pytest.approx(1e-6) + assert response.metrics["learning_rate"] == pytest.approx(2e-4) + + def test_optim_step_uses_registered_session_default_learning_rate(self): + """A native request can omit learning_rate when the session has an optimizer default.""" + server = APIServer( + engine_input_addr="tcp://127.0.0.1:17024", + engine_output_addr="tcp://127.0.0.1:17025", + ) + client = _FakeOrchestratorClient() + server._running = True + server.orchestrator_client = client + server.model_configs["default"] = { + "base_model": "Qwen/Qwen3-8B", + "optimizer_config": {"learning_rate": 7e-5}, + } + + async def _wait_for_response(self, response_future, request_id, timeout, timeout_message="timeout"): + return SimpleNamespace(outputs=[{"grad_norm": 0.5}]) + + server._wait_for_response = types.MethodType(_wait_for_response, server) + + response = asyncio.run(server.optim_step(OptimStepRequest(model_id="default"))) + + assert client.last_request.payload.lr == pytest.approx(7e-5) + assert response.metrics["learning_rate"] == pytest.approx(7e-5) + + def test_optim_step_uses_server_train_config_learning_rate_for_full_weight_default(self): + """Full-weight default sessions should inherit the server optimizer LR when request LR is omitted.""" + server = APIServer( + engine_input_addr="tcp://127.0.0.1:17034", + engine_output_addr="tcp://127.0.0.1:17035", + train_config={"lr": 6e-5}, + lora_config={"enable_lora": False}, + ) + client = _FakeOrchestratorClient() + server._running = True + server.orchestrator_client = client + server.model_configs["default"] = { + "base_model": "Qwen/Qwen3-8B", + "is_lora": False, + } + + async def _wait_for_response(self, response_future, request_id, timeout, timeout_message="timeout"): + return SimpleNamespace(outputs=[{"grad_norm": 0.5}]) + + server._wait_for_response = types.MethodType(_wait_for_response, server) + + response = asyncio.run(server.optim_step(OptimStepRequest(model_id="default"))) + + assert client.last_request.payload.lr == pytest.approx(6e-5) + assert response.metrics["learning_rate"] == pytest.approx(6e-5) + + def test_optim_step_uses_learning_rate_registered_by_create_model(self): + """create_model optimizer_config should feed later native optim_step requests.""" + default_session_spec, lora_config = _build_default_session_spec() + server = APIServer( + engine_input_addr="tcp://127.0.0.1:17028", + engine_output_addr="tcp://127.0.0.1:17029", + base_model="Qwen/Qwen3-8B", + default_session_spec=default_session_spec, + server_lora_config=lora_config, + max_lora_rank=16, + skip_initial_checkpoint=True, + ) + server.future_store = _ImmediateFutureStore() + client = _FakeOrchestratorClient() + server._running = True + server.orchestrator_client = client + + async def _wait_for_create_response(self, response_future, request_id, timeout, timeout_message="timeout"): + return SimpleNamespace( + outputs={"result": {"registered": True, "model_id": response_future.payload.model_id}} + ) + + server._wait_for_response = types.MethodType(_wait_for_create_response, server) + + asyncio.run( + create_model_endpoint( + CreateModelRequest( + model_id="session-from-create-model", + base_model="Qwen/Qwen3-8B", + optimizer_config={"learning_rate": 9e-5}, + ), + server=server, + ) + ) + + async def _wait_for_response(self, response_future, request_id, timeout, timeout_message="timeout"): + return SimpleNamespace(outputs=[{"grad_norm": 0.5}]) + + server._wait_for_response = types.MethodType(_wait_for_response, server) + + response = asyncio.run(server.optim_step(OptimStepRequest(model_id="session-from-create-model"))) + + assert client.last_request.payload.lr == pytest.approx(9e-5) + assert response.metrics["learning_rate"] == pytest.approx(9e-5) + + def test_optim_step_rejects_missing_learning_rate_without_session_default(self): + """Missing request and session learning rates should fail loudly instead of using a magic number.""" + server = APIServer( + engine_input_addr="tcp://127.0.0.1:17026", + engine_output_addr="tcp://127.0.0.1:17027", + ) + server._running = True + server.orchestrator_client = _FakeOrchestratorClient() + + with pytest.raises(HTTPException) as exc_info: + asyncio.run(server.optim_step(OptimStepRequest(model_id="default"))) + + assert exc_info.value.status_code == 400 + assert "no learning_rate" in exc_info.value.detail + + def test_save_weights_for_sampler_uses_lora_path_for_canonical_lora_config(self, tmp_path): + """Canonical lora_rank metadata should select LoRA-only sampler export.""" + server = APIServer( + engine_input_addr="tcp://127.0.0.1:17030", + engine_output_addr="tcp://127.0.0.1:17031", + output_dir=str(tmp_path), + ) + client = _FakeOrchestratorClient() + server._running = True + server.orchestrator_client = client + server.model_configs["session-a"] = { + "base_model": "Qwen/Qwen3-8B", + "lora_config": {"lora_rank": 8, "lora_alpha": 16}, + } + + async def _wait_for_response(self, response_future, request_id, timeout, timeout_message="timeout"): + return SimpleNamespace(outputs=[{"success": True, "lora_path": str(tmp_path / "adapter")}]) + + server._wait_for_response = types.MethodType(_wait_for_response, server) + + response = asyncio.run( + server.save_weights_for_sampler(SaveWeightsForSamplerRequest(model_id="session-a", name="step-1")) + ) + + assert response.path == "xorl://session-a/sampler_weights/step-1" + assert client.last_request.operation == "save_lora_only" + assert client.last_request.payload.model_id == "session-a" + 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 098b4ab2..013a89e5 100644 --- a/tests/server/api_server/test_api_types.py +++ b/tests/server/api_server/test_api_types.py @@ -10,6 +10,7 @@ from xorl.server.api_server.api_types import ( AdamParams, CreateModelRequest, + CreateSessionRequest, Datum, DatumInput, ErrorResponse, @@ -18,13 +19,16 @@ HealthCheckResponse, LoadWeightsRequest, LoadWeightsResponse, + LoRAConfigRequest, LossFnOutput, + OptimizerConfigRequest, OptimStepRequest, OptimStepResponse, SaveWeightsForSamplerRequest, SaveWeightsForSamplerResponse, SaveWeightsRequest, SaveWeightsResponse, + WeightsInfoResponse, ) @@ -137,29 +141,66 @@ class TestOptimWeightsHealthAndSerialization: """Test OptimStep, Weights, Health, Error types and serialization.""" def test_optim_step_types(self): - """Test AdamParams, OptimStepRequest and OptimStepResponse.""" - params = AdamParams(learning_rate=0.001, beta1=0.9, beta2=0.999, eps=1e-8) - assert params.learning_rate == 0.001 + """Test optimizer/session request types and OptimStepRequest/Response.""" + optimizer = OptimizerConfigRequest( + type="adamw", + learning_rate=0.001, + betas=[0.9, 0.999], + eps=1e-8, + ) + assert optimizer.learning_rate == 0.001 + assert optimizer.betas == [0.9, 0.999] - defaults = AdamParams() - assert defaults.learning_rate == 0.0001 - assert defaults.beta1 == 0.9 - assert defaults.beta2 == 0.95 - assert defaults.eps == 1e-12 + lora = LoRAConfigRequest(rank=8, alpha=16) + assert lora.lora_rank == 8 + assert lora.lora_alpha == 16 + assert lora.model_dump(exclude_none=True) == {"lora_rank": 8, "lora_alpha": 16} + assert set(LoRAConfigRequest.model_json_schema()["properties"]) == {"lora_rank", "lora_alpha"} + + create_request = CreateModelRequest( + model_id="session-a", + base_model="Qwen/Qwen3-8B", + lora_config=LoRAConfigRequest(rank=8, alpha=16), + optimizer_config=OptimizerConfigRequest(type="signsgd", learning_rate=2e-4), + ) + assert create_request.lora_config is not None + assert create_request.lora_config.lora_rank == 8 + assert create_request.optimizer_config is not None + assert create_request.optimizer_config.type == "signsgd" + + create_session_request = CreateSessionRequest(lora_config={"rank": 4, "alpha": 10}) + assert create_session_request.lora_config is not None + assert create_session_request.lora_config.model_dump(exclude_none=True) == { + "lora_rank": 4, + "lora_alpha": 10, + } request = OptimStepRequest( model_id="test-model", - adam_params=AdamParams(learning_rate=1e-4), + learning_rate=1e-4, gradient_clip=1.0, ) assert request.model_id == "test-model" - assert request.adam_params.learning_rate == 1e-4 + assert request.learning_rate == 1e-4 assert request.gradient_clip == 1.0 request = OptimStepRequest() assert request.model_id == "default" - assert request.adam_params.learning_rate == 0.0001 + assert request.learning_rate is None assert request.gradient_clip is None + assert request.adam_params is None + + legacy_request = OptimStepRequest( + **{ + "session_id": "legacy-session", + "adam_params": {"learning_rate": 2e-4, "grad_clip_norm": 1.5}, + } + ) + assert legacy_request.model_id == "legacy-session" + assert legacy_request.learning_rate == pytest.approx(2e-4) + assert legacy_request.gradient_clip == pytest.approx(1.5) + assert isinstance(legacy_request.adam_params, AdamParams) + assert legacy_request.adam_params.learning_rate == pytest.approx(2e-4) response = OptimStepResponse( metrics={"grad_norm": 1.234, "learning_rate": 1e-4, "step": 100}, @@ -169,6 +210,45 @@ def test_optim_step_types(self): response = OptimStepResponse(metrics={}, info={}) assert len(response.metrics) == 0 + weights_info = WeightsInfoResponse( + base_model="Qwen/Qwen3-8B", + is_lora=True, + lora_config={"lora_rank": 8, "lora_alpha": 16}, + optimizer_config={ + "type": "adamw", + "learning_rate": 1e-4, + "weight_decay": 0.01, + "optimizer_dtype": "bf16", + "betas": [0.9, 0.95], + "eps": 1e-8, + "optimizer_kwargs": {}, + }, + ) + assert weights_info.lora_config.lora_rank == 8 + assert weights_info.lora_rank == 8 + assert weights_info.optimizer_config.type == "adamw" + + legacy_weights_info = WeightsInfoResponse( + base_model="Qwen/Qwen3-8B", + is_lora=True, + lora_rank=8, + ) + assert legacy_weights_info.lora_rank == 8 + assert legacy_weights_info.lora_config is None + + full_weight_info = WeightsInfoResponse( + base_model="Qwen/Qwen3-8B", + is_lora=False, + ) + assert full_weight_info.lora_config is None + assert full_weight_info.optimizer_config is None + + partial_lora_info = WeightsInfoResponse( + base_model="Qwen/Qwen3-8B", + is_lora=True, + ) + assert partial_lora_info.lora_rank is None + def test_weights_health_error_and_serialization(self): """Test save/load/sampler types, health, error, and roundtrip serialization.""" # SaveWeightsRequest @@ -205,7 +285,9 @@ def test_weights_health_error_and_serialization(self): lora_config={"rank": 64}, ) assert request.model_id == "session-123" - assert request.lora_config["lora_rank"] == 64 + assert request.lora_config is not None + assert request.lora_config.lora_rank == 64 + assert "rank" not in request.lora_config.model_dump(exclude_none=True) # LoadWeightsResponse response = LoadWeightsResponse(path="xorl://default/weights/checkpoint-001") @@ -266,11 +348,23 @@ def test_weights_health_error_and_serialization(self): assert len(request2.forward_backward_input.data) == len(request.forward_backward_input.data) # OptimStepRequest - request = OptimStepRequest(adam_params=AdamParams(learning_rate=0.001), gradient_clip=1.0) + request = OptimStepRequest(learning_rate=0.001, gradient_clip=1.0) data = request.model_dump() - assert data["adam_params"]["learning_rate"] == 0.001 + assert data["learning_rate"] == 0.001 request2 = OptimStepRequest(**data) - assert request2.adam_params.learning_rate == request.adam_params.learning_rate + assert request2.learning_rate == request.learning_rate + + legacy_data = {"session_id": "legacy-session", "adam_params": {"learning_rate": 3e-4}} + request3 = OptimStepRequest(**legacy_data) + assert request3.model_id == "legacy-session" + assert request3.learning_rate == pytest.approx(3e-4) + + legacy_data_with_clip = { + "session_id": "legacy-session", + "adam_params": {"learning_rate": 3e-4, "grad_clip_norm": 0.75}, + } + request4 = OptimStepRequest(**legacy_data_with_clip) + assert request4.gradient_clip == pytest.approx(0.75) # HealthCheckResponse response = HealthCheckResponse( diff --git a/tests/server/api_server/test_checkpoint_paths.py b/tests/server/api_server/test_checkpoint_paths.py index d64dc5d2..6c301812 100644 --- a/tests/server/api_server/test_checkpoint_paths.py +++ b/tests/server/api_server/test_checkpoint_paths.py @@ -12,13 +12,16 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest -from fastapi import HTTPException from fastapi.exceptions import HTTPException from xorl.server.api_server.api_types import ( + CreateSamplingSessionRequest, DeleteCheckpointRequest, + InferenceEndpoint, ListCheckpointsRequest, LoadWeightsRequest, + RemoveInferenceEndpointRequest, + SaveWeightsForSamplerRequest, SaveWeightsRequest, ) from xorl.server.api_server.server import APIServer @@ -361,10 +364,12 @@ def test_sampler_listing_deletion_and_adapter_tracking(self): # --- Path resolution --- ap = os.path.join(self.temp_dir, "sampler_weights", "a-001") os.makedirs(ap, exist_ok=True) - name, path = self.server._resolve_model_path("sampler_weights/a-001") - assert name == "a-001" and path == ap - name, path = self.server._resolve_model_path("a-001") - assert name == "a-001" + model_id, name, path = self.server._resolve_model_path("sampler_weights/a-001") + assert model_id is None and name == "a-001" and path == ap + model_id, name, path = self.server._resolve_model_path("a-001") + assert model_id is None and name == "a-001" + model_id, name, path = self.server._resolve_model_path("xorl://session-a/sampler_weights/a-001") + assert model_id == "session-a" and name == "a-001" and path == ap with pytest.raises(HTTPException) as exc_info: self.server._resolve_model_path("nonexistent") @@ -378,6 +383,147 @@ def test_sampler_listing_deletion_and_adapter_tracking(self): assert self.server.loaded_sampling_loras["default"][-1] == ("adapter-001", "/path/1") assert len(self.server.loaded_sampling_loras) == 1 + def test_save_weights_for_sampler_uses_normalized_lora_session_spec(self): + """Normalized session specs should still export adapter-only sampler weights.""" + self.server._running = True + self.server.orchestrator_client = MagicMock() + self.server.orchestrator_client.send_request = AsyncMock(return_value=AsyncMock()) + self.server.model_configs["adapter-run"] = { + "base_model": "Qwen/Qwen3-8B", + "is_lora": True, + "lora_config": {"lora_rank": 8, "lora_alpha": 16}, + "optimizer_config": { + "type": "signsgd", + "learning_rate": 2e-4, + "weight_decay": 0.0, + "optimizer_dtype": "bf16", + "betas": None, + "eps": None, + "optimizer_kwargs": {}, + }, + } + + mock_output = MagicMock() + mock_output.outputs = [{"lora_path": os.path.join(self.temp_dir, "sampler_weights", "adapter-export")}] + + with patch.object(self.server, "_wait_for_response", return_value=mock_output): + response = asyncio.run( + self.server.save_weights_for_sampler( + SaveWeightsForSamplerRequest(model_id="adapter-run", name="adapter-export") + ) + ) + + request = self.server.orchestrator_client.send_request.await_args.args[0] + assert request.operation == "save_lora_only" + assert request.payload.lora_path.endswith("sampler_weights/adapter-export") + assert response.path == "xorl://adapter-run/sampler_weights/adapter-export" + + def test_create_sampling_session_prunes_stale_tracking_before_eviction(self): + """Stale tracked adapters should not force a bogus unload before loading a fresh adapter.""" + fresh_path = os.path.join(self.temp_dir, "sampler_weights", "fresh-001") + os.makedirs(fresh_path, exist_ok=True) + + self.server.inference_endpoints = [MagicMock()] + self.server.max_adapters_per_model = 1 + self.server.loaded_sampling_loras["default"] = [("stale-001", "/stale/path")] + self.server._get_loaded_adapters_from_endpoint = AsyncMock(return_value=[]) + self.server._load_lora_on_inference_endpoints = AsyncMock(return_value=True) + self.server._unload_lora_on_inference_endpoints = AsyncMock(return_value=True) + + response = asyncio.run( + self.server.create_sampling_session(CreateSamplingSessionRequest(model_path="fresh-001")) + ) + + assert response.lora_name == "fresh-001" + self.server._unload_lora_on_inference_endpoints.assert_not_awaited() + self.server._load_lora_on_inference_endpoints.assert_awaited_once_with("fresh-001", fresh_path) + assert self.server.loaded_sampling_loras["default"] == [("fresh-001", fresh_path)] + + def test_reconcile_preserves_tracking_when_loaded_adapter_query_fails(self): + """A transient endpoint introspection failure should not erase tracked sampler adapters.""" + self.server.inference_endpoints = [MagicMock()] + self.server.loaded_sampling_loras["default"] = [("tracked-001", "/tracked/path")] + self.server._get_loaded_adapters_from_endpoint = AsyncMock(return_value=None) + + loaded_by_endpoint = asyncio.run(self.server._reconcile_tracked_adapters("default")) + + assert loaded_by_endpoint == [set()] + assert self.server.loaded_sampling_loras["default"] == [("tracked-001", "/tracked/path")] + + def test_create_sampling_session_tracks_xorl_uri_under_uri_model_id(self): + """xorl:// sampler URIs should keep the embedded model_id for cleanup tracking.""" + uri_path = os.path.join(self.temp_dir, "sampler_weights", "session-a-adapter") + os.makedirs(uri_path, exist_ok=True) + + self.server.inference_endpoints = [MagicMock()] + self.server._get_loaded_adapters_from_endpoint = AsyncMock(return_value=[]) + self.server._load_lora_on_inference_endpoints = AsyncMock(return_value=True) + + response = asyncio.run( + self.server.create_sampling_session( + CreateSamplingSessionRequest(model_path="xorl://session-a/sampler_weights/session-a-adapter") + ) + ) + + assert response.lora_name == "session-a-adapter" + assert self.server.loaded_sampling_loras["session-a"] == [("session-a-adapter", uri_path)] + assert self.server.loaded_sampling_loras.get("default", []) == [] + + def test_create_sampling_session_tracks_plain_path_under_request_model_id(self): + """Explicit model_id should control tracking for plain sampler paths.""" + session_path = os.path.join(self.temp_dir, "sampler_weights", "shared-adapter") + os.makedirs(session_path, exist_ok=True) + + self.server.inference_endpoints = [MagicMock()] + self.server._get_loaded_adapters_from_endpoint = AsyncMock(return_value=[]) + self.server._load_lora_on_inference_endpoints = AsyncMock(return_value=True) + + response = asyncio.run( + self.server.create_sampling_session( + CreateSamplingSessionRequest(model_id="session-b", model_path="shared-adapter") + ) + ) + + assert response.lora_name == "shared-adapter" + assert self.server.loaded_sampling_loras["session-b"] == [("shared-adapter", session_path)] + assert self.server.loaded_sampling_loras.get("default", []) == [] + + def test_create_sampling_session_does_not_track_failed_loads(self): + """A failed sampling-session load should not leave a stale adapter in the tracking list.""" + failing_path = os.path.join(self.temp_dir, "sampler_weights", "failing-001") + os.makedirs(failing_path, exist_ok=True) + + self.server.inference_endpoints = [MagicMock()] + self.server._get_loaded_adapters_from_endpoint = AsyncMock(return_value=[]) + self.server._load_lora_on_inference_endpoints = AsyncMock( + side_effect=HTTPException(status_code=500, detail="load failed") + ) + + with pytest.raises(HTTPException): + asyncio.run(self.server.create_sampling_session(CreateSamplingSessionRequest(model_path="failing-001"))) + + assert self.server.loaded_sampling_loras.get("default", []) == [] + + def test_remove_last_inference_endpoint_clears_tracking(self): + """Dropping the last inference endpoint should also clear sampler tracking state.""" + self.server.inference_endpoints = [ + InferenceEndpoint( + host="sgl.local", + port=30000, + worker_port=29999, + world_size=1, + healthy=True, + server_info=None, + ) + ] + self.server.loaded_sampling_loras["default"] = [("adapter-001", "/path/1")] + + response = self.server.remove_inference_endpoint(RemoveInferenceEndpointRequest(host="sgl.local", port=30000)) + + assert response.success is True + assert self.server.inference_endpoints == [] + assert self.server.loaded_sampling_loras == {} + if __name__ == "__main__": pytest.main([__file__, "-v"]) diff --git a/tests/server/api_server/test_inference_endpoints.py b/tests/server/api_server/test_inference_endpoints.py index 602fff68..8d5daae9 100644 --- a/tests/server/api_server/test_inference_endpoints.py +++ b/tests/server/api_server/test_inference_endpoints.py @@ -79,8 +79,160 @@ def test_add_inference_endpoint_uses_single_endpoint_port(self, monkeypatch): assert response.success is True assert response.endpoint is not None assert response.endpoint.port == 30000 + assert response.endpoint.worker_port == 30000 assert "http://inference.example:29999/health" not in calls + def test_add_inference_endpoint_checks_explicit_worker_port(self, monkeypatch): + calls: list[str] = [] + responses = { + "http://inference.example:30000/health": FakeResponse(), + "http://inference.example:31000/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, worker_port=31000), + ) + ) + + assert response.success is True + assert response.endpoint is not None + assert response.endpoint.worker_port == 31000 + assert "http://inference.example:30000/health" in calls + assert "http://inference.example:31000/health" in calls + + def test_add_inference_endpoint_auto_sync_uses_detected_tp_size(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": 4}), + } + 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", + ) + server._running = True + server.orchestrator_client = MagicMock() + captured_endpoints = None + + async def fake_sync_weights_to_endpoints(endpoints, **_kwargs): + nonlocal captured_endpoints + captured_endpoints = endpoints + return {"success": True, "message": "ok", "transfer_time": 0.1, "total_bytes": 123} + + server._sync_weights_to_endpoints = fake_sync_weights_to_endpoints + + response = asyncio.run( + server.add_inference_endpoint( + AddInferenceEndpointRequest( + host="inference.example", + port=30000, + world_size=1, + sync_weights=True, + master_address="train.example", + ), + ) + ) + + assert response.success is True + assert response.endpoint is not None + assert response.endpoint.world_size == 4 + assert captured_endpoints == [{"host": "inference.example", "port": 30000, "world_size": 4}] + + def test_list_inference_endpoints_accepts_v1_models_health_fallback(self, monkeypatch): + calls: list[str] = [] + responses = { + "http://inference.example:30000/health": FakeResponse(status_code=404), + "http://inference.example:30000/v1/models": FakeResponse(json_data={"data": []}), + } + 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", + ) + server.inference_endpoints = [ + InferenceEndpoint(host="inference.example", port=30000, world_size=1), + ] + + response = asyncio.run(server.list_inference_endpoints()) + + assert response.count == 1 + assert response.endpoints[0].host == "inference.example" + assert "http://inference.example:30000/health" in calls + assert "http://inference.example:30000/v1/models" in calls + + def test_lora_adapter_management_uses_worker_port(self, monkeypatch): + calls: list[str] = [] + + def fake_post(url: str, **kwargs): + calls.append(url) + return FakeResponse(json_data={"success": True}) + + monkeypatch.setattr("xorl.server.api_server.inference_endpoints.requests.post", fake_post) + + server = APIServer( + engine_input_addr="tcp://127.0.0.1:17002", + engine_output_addr="tcp://127.0.0.1:17003", + ) + server.inference_endpoints = [ + InferenceEndpoint(host="inference.example", port=30000, worker_port=31000, world_size=1), + ] + + asyncio.run(server._load_lora_on_inference_endpoints("adapter-001", "/tmp/adapter-001")) + asyncio.run(server._unload_lora_on_inference_endpoints("adapter-001")) + + assert calls == [ + "http://inference.example:31000/load_lora_adapter", + "http://inference.example:31000/unload_lora_adapter", + ] + + def test_loaded_adapter_query_uses_worker_port(self, monkeypatch): + calls: list[str] = [] + responses = { + "http://inference.example:31000/v1/models": FakeResponse( + json_data={ + "data": [ + {"id": "base-model"}, + {"id": "adapter-001", "parent": "base-model"}, + ] + } + ), + } + 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", + ) + endpoint = InferenceEndpoint(host="inference.example", port=30000, worker_port=31000, world_size=1) + + adapters = asyncio.run(server._get_loaded_adapters_from_endpoint(endpoint)) + + assert adapters == ["adapter-001"] + assert calls == ["http://inference.example:31000/v1/models"] + def test_sync_inference_weights_forwards_single_endpoint(self): server = APIServer( engine_input_addr="tcp://127.0.0.1:17002", @@ -107,6 +259,8 @@ async def fake_send_request(engine_request): "total_bytes": 123, "num_parameters": 4, "num_buckets": 1, + "timing_breakdown": {"transfer_s": 0.1, "total_handler_s": 0.2}, + "p2p_rank_summaries": [{"rank": 0, "is_sender": True, "transfer_wall_s": 0.1}], "endpoint_results": [{"host": "inference.example", "port": 30000, "success": True}], } ], @@ -130,3 +284,59 @@ async def fake_send_request(engine_request): "world_size": 2, } ] + assert response.timing_breakdown == {"transfer_s": 0.1, "total_handler_s": 0.2} + assert response.p2p_rank_summaries == [{"rank": 0, "is_sender": True, "transfer_wall_s": 0.1}] + + def test_add_inference_endpoint_auto_sync_uses_configured_sync_method(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", + sync_inference_method="p2p", + ) + server._running = True + 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, + } + ], + ) + ) + return future + + server.orchestrator_client = MagicMock(send_request=AsyncMock(side_effect=fake_send_request)) + + response = asyncio.run( + server.add_inference_endpoint( + AddInferenceEndpointRequest( + host="inference.example", + port=30000, + sync_weights=True, + master_address="train.example", + ), + ) + ) + + assert response.success is True + assert response.weights_synced is True + assert captured_request["request"].payload.sync_method == "p2p" diff --git a/tests/server/api_server/test_session_endpoints.py b/tests/server/api_server/test_session_endpoints.py new file mode 100644 index 00000000..b39f1da7 --- /dev/null +++ b/tests/server/api_server/test_session_endpoints.py @@ -0,0 +1,634 @@ +"""Focused tests for session-spec API endpoint behavior.""" + +from __future__ import annotations + +import json +import os +import types +from unittest.mock import AsyncMock + +import pytest +from fastapi import HTTPException + +from xorl.server.api_server.api_types import ( + CreateModelRequest, + CreateSessionRequest, + KillSessionRequest, + OptimizerConfigRequest, + SaveWeightsForSamplerRequest, + SaveWeightsResponse, + UnloadModelRequest, + WeightsInfoRequest, +) +from xorl.server.api_server.endpoints import ( + create_model_endpoint, + create_session_endpoint, + kill_session_endpoint, + unload_model_endpoint, + weights_info_endpoint, +) +from xorl.server.api_server.server import APIServer +from xorl.server.session_spec import build_default_session_spec, load_session_spec_from_checkpoint, write_session_spec + + +pytestmark = [pytest.mark.cpu, pytest.mark.server, pytest.mark.anyio] + + +class _ImmediateFutureStore: + def __init__(self) -> None: + self.last_result = None + self.deleted_models = [] + + async def create(self, *, model_id, request_type, process_fn, request_data, ttl=None): + self.last_result = await process_fn(request_data) + return "future-test-1" + + async def delete_by_model(self, model_id): + self.deleted_models.append(model_id) + return 0 + + +class _FakeOrchestratorClient: + def __init__(self) -> None: + self.last_request = None + self.requests = [] + + async def send_request(self, request): + self.last_request = request + self.requests.append(request) + return request + + +def _build_wait_for_response(): + async def _wait_for_response(self, response_future, request_id, timeout, timeout_message="timeout"): + if response_future.operation == "register_session": + return types.SimpleNamespace( + outputs={"result": {"registered": True, "model_id": response_future.payload.model_id}} + ) + if response_future.operation == "save_state": + os.makedirs(response_future.payload.checkpoint_path, exist_ok=True) + return types.SimpleNamespace(outputs=[{"checkpoint_path": response_future.payload.checkpoint_path}]) + if response_future.operation == "save_lora_only": + os.makedirs(response_future.payload.lora_path, exist_ok=True) + return types.SimpleNamespace(outputs=[{"lora_path": response_future.payload.lora_path}]) + if response_future.operation == "save_full_weights": + os.makedirs(response_future.payload.output_path, exist_ok=True) + return types.SimpleNamespace(outputs=[{"output_path": response_future.payload.output_path}]) + if response_future.operation == "kill_session": + return types.SimpleNamespace( + outputs=[ + { + "success": True, + "message": f"killed {response_future.payload.model_id}", + "checkpoint_path": None, + } + ] + ) + raise AssertionError(f"Unexpected operation: {response_future.operation}") + + return _wait_for_response + + +def _build_server(tmp_path): + train_config = { + "optimizer": "adamw", + "lr": 1e-5, + "weight_decay": 0.01, + "optimizer_dtype": "bf16", + "optimizer_kwargs": {}, + } + lora_config = { + "enable_lora": True, + "lora_rank": 8, + "max_lora_rank": 16, + "lora_alpha": 16, + "lora_target_modules": ["q_proj", "o_proj"], + } + default_session_spec = build_default_session_spec( + base_model="Qwen/Qwen3-8B", + train_config=train_config, + lora_config=lora_config, + ) + server = APIServer( + engine_input_addr="tcp://127.0.0.1:17000", + engine_output_addr="tcp://127.0.0.1:17001", + output_dir=str(tmp_path), + base_model="Qwen/Qwen3-8B", + default_session_spec=default_session_spec, + server_lora_config=lora_config, + max_lora_rank=16, + ) + server.future_store = _ImmediateFutureStore() + server.orchestrator_client = _FakeOrchestratorClient() + server._running = True + server._wait_for_response = types.MethodType(_build_wait_for_response(), server) + return server + + +def _build_full_weight_server(tmp_path): + server = APIServer( + engine_input_addr="tcp://127.0.0.1:17000", + engine_output_addr="tcp://127.0.0.1:17001", + output_dir=str(tmp_path), + base_model="Qwen/Qwen3-8B", + default_session_spec=None, + ) + server.future_store = _ImmediateFutureStore() + server.orchestrator_client = _FakeOrchestratorClient() + server._running = True + server._wait_for_response = types.MethodType(_build_wait_for_response(), server) + return server + + +async def test_create_model_endpoint_registers_normalized_session_spec(tmp_path): + server = _build_server(tmp_path) + + response = await create_model_endpoint( + CreateModelRequest( + model_id="session-a", + base_model="Qwen/Qwen3-8B", + lora_config={"rank": 4, "alpha": 12}, + optimizer_config=OptimizerConfigRequest(type="signsgd", learning_rate=2e-4), + ), + server=server, + ) + + assert response.request_id == "future-test-1" + assert "session-a" in server.model_configs + assert server.model_configs["session-a"]["lora_config"]["lora_rank"] == 4 + assert server.model_configs["session-a"]["lora_config"]["lora_alpha"] == 12 + assert server.model_configs["session-a"]["optimizer_config"]["type"] == "signsgd" + register_requests = [ + request for request in server.orchestrator_client.requests if request.operation == "register_session" + ] + assert len(register_requests) == 1 + assert register_requests[0].payload.session_spec["lora_config"]["lora_rank"] == 4 + + +async def test_create_session_endpoint_registers_lora_session_with_workers(tmp_path): + server = _build_server(tmp_path) + + response = await create_session_endpoint( + CreateSessionRequest(session_id="session-a", lora_config={"rank": 4, "alpha": 12}), + server=server, + ) + + assert response.session_id == "session-a" + assert "session-a" in server.registered_model_ids + assert server.model_configs["session-a"]["lora_config"] == {"lora_rank": 4, "lora_alpha": 12} + register_requests = [ + request for request in server.orchestrator_client.requests if request.operation == "register_session" + ] + assert len(register_requests) == 1 + assert register_requests[0].payload.model_id == "session-a" + assert register_requests[0].payload.materialize is True + assert register_requests[0].payload.session_spec["lora_config"]["lora_rank"] == 4 + + +async def test_create_session_endpoint_allows_rank_only_lora_override(tmp_path): + server = _build_server(tmp_path) + + response = await create_session_endpoint( + CreateSessionRequest(session_id="session-rank-only", lora_config={"rank": 4}), + server=server, + ) + + assert response.session_id == "session-rank-only" + assert server.model_configs["session-rank-only"]["lora_config"] == {"lora_rank": 4, "lora_alpha": 16} + register_requests = [ + request for request in server.orchestrator_client.requests if request.operation == "register_session" + ] + assert len(register_requests) == 1 + assert register_requests[0].payload.session_spec["lora_config"] == {"lora_rank": 4, "lora_alpha": 16} + + +async def test_create_session_endpoint_refreshes_existing_custom_lora_session(tmp_path): + server = _build_server(tmp_path) + + await create_session_endpoint( + CreateSessionRequest(session_id="session-a", lora_config={"rank": 4, "alpha": 12}), + server=server, + ) + + response = await create_session_endpoint( + CreateSessionRequest(session_id="session-a"), + server=server, + ) + + assert response.session_id == "session-a" + assert response.warning_message == "Session 'session-a' already existed; refreshed activity timestamp." + assert server.model_configs["session-a"]["lora_config"] == {"lora_rank": 4, "lora_alpha": 12} + register_requests = [ + request for request in server.orchestrator_client.requests if request.operation == "register_session" + ] + assert len(register_requests) == 1 + + +async def test_create_model_endpoint_rejects_conflicting_recreate(tmp_path): + server = _build_server(tmp_path) + + await create_model_endpoint( + CreateModelRequest( + model_id="session-a", + base_model="Qwen/Qwen3-8B", + lora_config={"rank": 4}, + ), + server=server, + ) + + with pytest.raises(ValueError, match="already exists with a different session spec"): + await create_model_endpoint( + CreateModelRequest( + model_id="session-a", + base_model="Qwen/Qwen3-8B", + lora_config={"rank": 6}, + ), + server=server, + ) + + +async def test_create_model_endpoint_propagates_register_session_failure(tmp_path): + server = _build_server(tmp_path) + + async def _fail_wait_for_response(self, response_future, request_id, timeout, timeout_message="timeout"): + raise HTTPException(status_code=500, detail="Engine error: Cross-rank error: rank 1: register failed") + + server._wait_for_response = types.MethodType(_fail_wait_for_response, server) + + with pytest.raises(HTTPException, match="Cross-rank error"): + await create_model_endpoint( + CreateModelRequest( + model_id="session-b", + base_model="Qwen/Qwen3-8B", + lora_config={"rank": 4}, + ), + server=server, + ) + + assert "session-b" not in server.model_configs + assert "session-b" not in server.registered_model_ids + + +async def test_create_model_endpoint_ensures_reserved_checkpoint_for_default_session(tmp_path): + server = _build_server(tmp_path) + server.save_weights = AsyncMock(return_value=SaveWeightsResponse(path="xorl://default/weights/000000")) + + response = await create_model_endpoint( + CreateModelRequest( + model_id="default", + base_model="Qwen/Qwen3-8B", + ), + server=server, + ) + + assert response.request_id == "future-test-1" + assert server.save_weights.await_count == 1 + save_request = server.save_weights.await_args.args[0] + assert save_request.model_id == "default" + assert save_request.path == "000000" + + +async def test_create_model_endpoint_saves_reserved_checkpoint_per_session(tmp_path): + server = _build_server(tmp_path) + server.save_weights = AsyncMock( + side_effect=[ + SaveWeightsResponse(path="xorl://session-a/weights/000000"), + SaveWeightsResponse(path="xorl://session-b/weights/000000"), + ] + ) + + await create_model_endpoint( + CreateModelRequest( + model_id="session-a", + base_model="Qwen/Qwen3-8B", + lora_config={"rank": 4}, + ), + server=server, + ) + await create_model_endpoint( + CreateModelRequest( + model_id="session-b", + base_model="Qwen/Qwen3-8B", + lora_config={"rank": 6}, + ), + server=server, + ) + + assert server.save_weights.await_count == 2 + assert [call.args[0].model_id for call in server.save_weights.await_args_list] == ["session-a", "session-b"] + assert [call.args[0].path for call in server.save_weights.await_args_list] == ["000000", "000000"] + + +async def test_create_model_endpoint_overwrites_stale_reserved_checkpoint_for_recreated_session(tmp_path): + server = _build_server(tmp_path) + checkpoint_dir = tmp_path / "weights" / "session-a" / server.RESERVED_CHECKPOINT_NAME + checkpoint_dir.mkdir(parents=True) + server.save_weights = AsyncMock(return_value=SaveWeightsResponse(path="xorl://session-a/weights/000000")) + + await create_model_endpoint( + CreateModelRequest( + model_id="session-a", + base_model="Qwen/Qwen3-8B", + lora_config={"rank": 4}, + ), + server=server, + ) + + assert server.save_weights.await_count == 1 + save_request = server.save_weights.await_args.args[0] + assert save_request.model_id == "session-a" + assert save_request.path == "000000" + + +async def test_create_model_endpoint_preserves_reserved_checkpoint_for_existing_session(tmp_path): + server = _build_server(tmp_path) + server.save_weights = AsyncMock(return_value=SaveWeightsResponse(path="xorl://session-a/weights/000000")) + + request = CreateModelRequest( + model_id="session-a", + base_model="Qwen/Qwen3-8B", + lora_config={"rank": 4}, + ) + await create_model_endpoint(request, server=server) + (tmp_path / "weights" / "session-a" / server.RESERVED_CHECKPOINT_NAME).mkdir(parents=True) + + await create_model_endpoint(request, server=server) + + assert server.save_weights.await_count == 1 + + +async def test_save_weights_for_sampler_uses_lora_only_for_normalized_session(tmp_path): + server = _build_server(tmp_path) + + response = await server.save_weights_for_sampler( + SaveWeightsForSamplerRequest(model_id="default", name="sampler-a"), + ) + + assert response.path == "xorl://default/sampler_weights/sampler-a" + assert server.orchestrator_client.requests[-1].operation == "save_lora_only" + assert server.orchestrator_client.requests[-1].payload.model_id == "default" + + +async def test_create_model_endpoint_rejects_full_weight_multitenancy(tmp_path): + server = _build_full_weight_server(tmp_path) + + with pytest.raises(ValueError, match="multi-tenancy is not supported yet"): + await create_model_endpoint( + CreateModelRequest( + model_id="session-a", + base_model="Qwen/Qwen3-8B", + ), + server=server, + ) + + +async def test_create_model_endpoint_allows_default_full_weight_session(tmp_path): + server = _build_full_weight_server(tmp_path) + + response = await create_model_endpoint( + CreateModelRequest( + model_id="default", + base_model="Qwen/Qwen3-8B", + ), + server=server, + ) + + assert response.request_id == "future-test-1" + assert server.model_configs["default"] == { + "base_model": "Qwen/Qwen3-8B", + "is_lora": False, + } + register_requests = [ + request for request in server.orchestrator_client.requests if request.operation == "register_session" + ] + assert len(register_requests) == 1 + assert register_requests[0].payload.materialize is False + + +async def test_create_model_endpoint_rejects_full_weight_overrides(tmp_path): + server = _build_full_weight_server(tmp_path) + + with pytest.raises(ValueError, match="Per-session LoRA or optimizer overrides are not supported"): + await create_model_endpoint( + CreateModelRequest( + model_id="default", + base_model="Qwen/Qwen3-8B", + optimizer_config=OptimizerConfigRequest(type="signsgd", learning_rate=2e-4), + ), + server=server, + ) + + +async def test_kill_session_endpoint_cleans_up_lora_session_registry(tmp_path): + server = _build_server(tmp_path) + request = CreateModelRequest( + model_id="session-a", + base_model="Qwen/Qwen3-8B", + lora_config={"rank": 4}, + ) + + await create_model_endpoint( + request, + server=server, + ) + + response = await kill_session_endpoint( + KillSessionRequest(model_id="session-a", save_checkpoint=False), + server=server, + ) + + assert response.success is True + assert "session-a" not in server.registered_model_ids + assert "session-a" not in server.model_configs + assert server.future_store.deleted_models == ["session-a"] + kill_requests = [request for request in server.orchestrator_client.requests if request.operation == "kill_session"] + assert len(kill_requests) == 1 + assert kill_requests[0].payload.model_id == "session-a" + + await create_model_endpoint(request, server=server) + register_requests = [ + request for request in server.orchestrator_client.requests if request.operation == "register_session" + ] + assert len(register_requests) == 2 + + +async def test_kill_session_endpoint_returns_xorl_uri_for_lora_checkpoint(tmp_path): + server = _build_server(tmp_path) + request = CreateModelRequest( + model_id="session-a", + base_model="Qwen/Qwen3-8B", + lora_config={"rank": 4}, + ) + await create_model_endpoint(request, server=server) + + checkpoint_dir = tmp_path / "weights" / "session-a" / "session_session-a_final" + checkpoint_dir.mkdir(parents=True) + + async def _wait_for_kill(self, response_future, request_id, timeout, timeout_message="timeout"): + if response_future.operation == "kill_session": + return types.SimpleNamespace( + outputs=[ + { + "success": True, + "message": "killed session-a", + "checkpoint_path": str(checkpoint_dir), + } + ] + ) + return await _build_wait_for_response()(self, response_future, request_id, timeout, timeout_message) + + server._wait_for_response = types.MethodType(_wait_for_kill, server) + + response = await kill_session_endpoint( + KillSessionRequest(model_id="session-a", save_checkpoint=True), + server=server, + ) + + assert response.success is True + assert response.checkpoint_path == "xorl://session-a/weights/session_session-a_final" + + +async def test_kill_session_endpoint_preserves_default_lora_session(tmp_path): + server = _build_server(tmp_path) + + response = await kill_session_endpoint( + KillSessionRequest(model_id="default", save_checkpoint=False), + server=server, + ) + + assert response.success is True + assert "default" in server.registered_model_ids + assert "default" in server.model_configs + assert server.orchestrator_client.requests == [] + + +async def test_unload_model_endpoint_rejects_default_lora_session(tmp_path): + server = _build_server(tmp_path) + + with pytest.raises(HTTPException, match="reserved and cannot be unloaded") as exc_info: + await unload_model_endpoint( + UnloadModelRequest(model_id="default"), + server=server, + ) + + assert exc_info.value.status_code == 400 + + +async def test_weights_info_endpoint_reads_session_spec_from_checkpoint(tmp_path): + server = _build_server(tmp_path) + checkpoint_dir = tmp_path / "weights" / "session-a" / "ckpt-001" + checkpoint_dir.mkdir(parents=True) + + write_session_spec( + str(checkpoint_dir), + { + "base_model": "Qwen/Qwen3-8B", + "is_lora": True, + "lora_config": {"lora_rank": 4, "lora_alpha": 12}, + "optimizer_config": { + "type": "signsgd", + "learning_rate": 2e-4, + "weight_decay": 0.0, + "optimizer_dtype": "bf16", + "betas": None, + "eps": None, + "optimizer_kwargs": {}, + }, + }, + ) + + # Intentionally store different in-memory metadata to ensure weights_info trusts disk. + server.model_configs["session-a"] = { + "base_model": "Qwen/Qwen3-8B", + "is_lora": True, + "lora_config": {"lora_rank": 8, "lora_alpha": 16}, + "optimizer_config": { + "type": "adamw", + "learning_rate": 1e-5, + "weight_decay": 0.01, + "optimizer_dtype": "bf16", + "betas": [0.9, 0.95], + "eps": 1e-8, + "optimizer_kwargs": {}, + }, + } + + response = await weights_info_endpoint( + WeightsInfoRequest(xorl_path="xorl://session-a/weights/ckpt-001"), + server=server, + ) + + assert response.base_model == "Qwen/Qwen3-8B" + assert response.lora_config.lora_rank == 4 + assert response.lora_config.lora_alpha == 12 + assert response.optimizer_config.type == "signsgd" + assert response.optimizer_config.learning_rate == pytest.approx(2e-4) + + +async def test_weights_info_endpoint_rejects_checkpoint_path_escape(tmp_path): + server = _build_server(tmp_path / "output") + escaped_dir = tmp_path / "secret-checkpoint" + escaped_dir.mkdir() + write_session_spec(str(escaped_dir), {"base_model": "escaped", "is_lora": False}) + (tmp_path / "output" / "weights" / "default").mkdir(parents=True) + + with pytest.raises(HTTPException) as exc_info: + await weights_info_endpoint( + WeightsInfoRequest(xorl_path="xorl://default/weights/../../../secret-checkpoint"), + server=server, + ) + + assert exc_info.value.status_code == 400 + + +async def test_weights_info_endpoint_returns_full_weight_checkpoint_metadata(tmp_path): + server = _build_full_weight_server(tmp_path) + checkpoint_dir = tmp_path / "weights" / "default" / "ckpt-001" + checkpoint_dir.mkdir(parents=True) + server.model_configs["default"] = { + "base_model": "Qwen/Qwen3-8B", + "is_lora": False, + } + + response = await weights_info_endpoint( + WeightsInfoRequest(xorl_path="xorl://default/weights/ckpt-001"), + server=server, + ) + + assert response.base_model == "Qwen/Qwen3-8B" + assert response.is_lora is False + assert response.lora_config is None + assert response.optimizer_config is None + + +def test_load_session_spec_from_checkpoint_upgrades_legacy_signsgd_metadata(tmp_path): + checkpoint_dir = tmp_path / "weights" / "session-a" / "ckpt-legacy" + checkpoint_dir.mkdir(parents=True) + + metadata = { + "lr": 2e-4, + "optimizer": { + "type": "signsgd", + "weight_decay": 0.0, + "betas": None, + "eps": None, + "optimizer_kwargs": {}, + }, + } + adapter_config = { + "base_model_name_or_path": "Qwen/Qwen3-8B", + "r": 4, + "lora_alpha": 12, + } + (checkpoint_dir / "metadata.json").write_text(json.dumps(metadata), encoding="utf-8") + (checkpoint_dir / "adapter_config.json").write_text(json.dumps(adapter_config), encoding="utf-8") + + session_spec = load_session_spec_from_checkpoint(str(checkpoint_dir)) + + assert session_spec["base_model"] == "Qwen/Qwen3-8B" + assert session_spec["is_lora"] is True + assert session_spec["lora_config"]["lora_rank"] == 4 + assert session_spec["lora_config"]["lora_alpha"] == 12 + assert session_spec["optimizer_config"]["type"] == "signsgd" + assert session_spec["optimizer_config"]["betas"] is None + assert session_spec["optimizer_config"]["eps"] is None diff --git a/tests/server/api_server/test_training_ops.py b/tests/server/api_server/test_training_ops.py new file mode 100644 index 00000000..28acea8e --- /dev/null +++ b/tests/server/api_server/test_training_ops.py @@ -0,0 +1,114 @@ +"""Focused tests for API training operation responses.""" + +from __future__ import annotations + +import types + +import pytest + +from xorl.server.api_server.api_types import ForwardRequest, OptimStepRequest +from xorl.server.api_server.server import APIServer + + +pytestmark = [pytest.mark.cpu, pytest.mark.server, pytest.mark.anyio] + + +class _FakeOrchestratorClient: + def __init__(self) -> None: + self.last_request = None + + async def send_request(self, request): + self.last_request = request + return request + + +def _build_wait_for_response(): + async def _wait_for_response(self, response_future, request_id, timeout, timeout_message="timeout"): + return types.SimpleNamespace( + outputs=[ + { + "grad_norm": 7.5, + "learning_rate": 2e-4, + "step": 1, + } + ] + ) + + return _wait_for_response + + +def _build_server(): + server = APIServer( + engine_input_addr="tcp://127.0.0.1:17000", + engine_output_addr="tcp://127.0.0.1:17001", + ) + server.orchestrator_client = _FakeOrchestratorClient() + server._running = True + server._wait_for_response = types.MethodType(_build_wait_for_response(), server) + return server + + +async def test_optim_step_uses_orchestrator_learning_rate_key(): + server = _build_server() + + response = await server.optim_step(OptimStepRequest(model_id="test-session", learning_rate=2e-4, gradient_clip=1.0)) + + assert response.metrics["grad_norm"] == pytest.approx(7.5) + assert response.metrics["learning_rate"] == pytest.approx(2e-4) + assert server.orchestrator_client.last_request.payload.lr == pytest.approx(2e-4) + + +async def test_optim_step_maps_legacy_grad_clip_norm_to_orchestrator_payload(): + server = _build_server() + + response = await server.optim_step( + OptimStepRequest( + **{ + "session_id": "legacy-session", + "adam_params": {"learning_rate": 3e-4, "grad_clip_norm": 2.5}, + } + ) + ) + + assert response.metrics["grad_norm"] == pytest.approx(7.5) + assert server.orchestrator_client.last_request.payload.lr == pytest.approx(3e-4) + assert server.orchestrator_client.last_request.payload.gradient_clip == pytest.approx(2.5) + + +async def test_forward_surfaces_auto_load_info(): + server = _build_server() + + async def _wait_for_response(self, response_future, request_id, timeout, timeout_message="timeout"): + return types.SimpleNamespace( + outputs=[ + { + "loss": 0.25, + "valid_tokens": 2, + "execution_time": 0.01, + "auto_loaded": True, + "auto_load_path": "/tmp/evicted/session-a", + } + ] + ) + + server._wait_for_response = types.MethodType(_wait_for_response, server) + + response = await server.forward( + ForwardRequest( + model_id="session-a", + forward_input={ + "data": [ + { + "model_input": {"input_ids": [1, 2]}, + "loss_fn_inputs": {"labels": [1, 2]}, + } + ] + }, + ) + ) + + assert response.metrics["loss:mean"] == pytest.approx(0.25) + assert response.info == { + "auto_loaded": True, + "auto_load_path": "/tmp/evicted/session-a", + } diff --git a/tests/server/backend/test_remote_backend.py b/tests/server/backend/test_remote_backend.py new file mode 100644 index 00000000..1bcdb65c --- /dev/null +++ b/tests/server/backend/test_remote_backend.py @@ -0,0 +1,28 @@ +import pytest + +from xorl.server.backend.remote import RemoteBackend + + +@pytest.mark.asyncio +async def test_sync_inference_weights_uses_backend_operation_timeout(monkeypatch): + backend = RemoteBackend(operation_timeout=2400.0) + captured = {} + + async def fake_execute(operation, payload, request_id=None, timeout=None): + captured["operation"] = operation + captured["payload"] = payload + captured["request_id"] = request_id + captured["timeout"] = timeout + return {"success": True} + + monkeypatch.setattr(backend, "_execute", fake_execute) + + await backend.sync_inference_weights( + endpoints=[{"host": "inference.example", "port": 30000, "world_size": 4}], + master_address="trainer.example", + request_id="sync-req", + ) + + assert captured["operation"] == "sync_inference_weights" + assert captured["request_id"] == "sync-req" + assert captured["timeout"] == 2400.0 diff --git a/tests/server/orchestrator/test_orchestrator.py b/tests/server/orchestrator/test_orchestrator.py index f2f5ddb3..3bd2526a 100644 --- a/tests/server/orchestrator/test_orchestrator.py +++ b/tests/server/orchestrator/test_orchestrator.py @@ -13,14 +13,11 @@ - Tests focus on Orchestrator's scheduling, routing, and output formatting """ -import msgpack -import pytest - - -pytestmark = [pytest.mark.cpu, pytest.mark.server] import socket import time +import msgpack +import pytest import zmq from xorl.server.backend import DummyBackend @@ -39,6 +36,9 @@ ) +pytestmark = [pytest.mark.cpu, pytest.mark.server] + + # ============================================================================ # Fixtures # ============================================================================ @@ -93,7 +93,7 @@ def output_socket(addresses): context = zmq.Context() sock = context.socket(zmq.PULL) sock.setsockopt(zmq.LINGER, 0) - sock.connect(addresses["output"]) + sock.bind(addresses["output"]) yield sock diff --git a/tests/server/orchestrator/test_orchestrator_client_communication.py b/tests/server/orchestrator/test_orchestrator_client_communication.py index 29b71413..7cedb9b5 100644 --- a/tests/server/orchestrator/test_orchestrator_client_communication.py +++ b/tests/server/orchestrator/test_orchestrator_client_communication.py @@ -43,7 +43,7 @@ class MockEngine: This mock implements the same socket pattern as the real engine: - INPUT socket (DEALER): connects to API server's ROUTER to receive requests - - OUTPUT socket (PUSH): binds for API server's PULL to receive outputs + - OUTPUT socket (PUSH): connects to API server's PULL to send outputs """ def __init__(self, input_addr, output_addr, engine_identity=b"engine-0", response_delay=0.0): @@ -69,7 +69,7 @@ async def start(self): self.input_socket.connect(self.input_addr) self.output_socket = self.context.socket(zmq.PUSH) self.output_socket.setsockopt(zmq.LINGER, 0) - self.output_socket.bind(self.output_addr) + self.output_socket.connect(self.output_addr) await asyncio.sleep(0.2) self._running = True self._task = asyncio.create_task(self._process_requests()) diff --git a/tests/server/orchestrator/test_packing.py b/tests/server/orchestrator/test_packing.py index 7fe6c9c2..beb00447 100644 --- a/tests/server/orchestrator/test_packing.py +++ b/tests/server/orchestrator/test_packing.py @@ -4,9 +4,6 @@ import pytest import torch - -pytestmark = [pytest.mark.cpu, pytest.mark.server] - from xorl.server.orchestrator.packing import ( Packer, SequentialPacker, @@ -16,6 +13,9 @@ ) +pytestmark = [pytest.mark.cpu, pytest.mark.server] + + # ============================================================================ # Fixtures # ============================================================================ @@ -63,6 +63,80 @@ def test_packing_enabled(simple_data): assert batch["position_ids"] == [[0, 1, 2, 0, 0, 1]] +def test_opd_metadata_packs_as_token_aligned_fields(): + """OPD teacher ids, cache refs, and weights survive packed dispatch.""" + data = [ + { + "input_ids": [1, 2, 3], + "target_tokens": [2, 3, 4], + "teacher_ids": [0, 0, 0], + "teacher_cache_indices": [10, 11, 12], + "teacher_weights": [1.0, 1.0, 0.5], + }, + { + "input_ids": [5, 6], + "target_tokens": [6, 7], + "teacher_id": 2, + "teacher_cache_indices": [20, 21], + "teacher_weight": 0.25, + }, + ] + packer = SequentialPacker(enable_packing=True, log_stats=False, pad_to_multiple_of=1) + batches = packer.pack(data, max_seq_len=100, request_id="opd") + + assert len(batches) == 1 + batch = batches[0] + assert batch["labels"] == [[2, 3, 4, 6, 7]] + assert batch["teacher_ids"] == [[0, 0, 0, 2, 2]] + assert batch["teacher_cache_indices"] == [[10, 11, 12, 20, 21]] + assert batch["teacher_weights"] == [[1.0, 1.0, 0.5, 0.25, 0.25]] + + +def test_opd_metadata_shifts_with_hf_style_labels(): + """OPD per-token fields stay aligned when packing shifts HF-style labels.""" + data = [ + { + "input_ids": [1, 2, 3, 4], + "labels": [10, 20, 30, 40], + "teacher_ids": [0, 0, 1, 1], + "teacher_cache_indices": [100, 101, 102, 103], + "teacher_weights": [1.0, 0.5, 0.25, 0.125], + "teacher_hidden_states": [ + [1.0, 1.5], + [2.0, 2.5], + [3.0, 3.5], + [4.0, 4.5], + ], + } + ] + packer = SequentialPacker(enable_packing=True, log_stats=False, pad_to_multiple_of=1) + batches = packer.pack(data, max_seq_len=100, request_id="opd-hf") + + assert len(batches) == 1 + batch = batches[0] + assert batch["input_ids"] == [[1, 2, 3]] + assert batch["labels"] == [[20, 30, 40]] + assert batch["teacher_ids"] == [[0, 0, 1]] + assert batch["teacher_cache_indices"] == [[100, 101, 102]] + assert batch["teacher_weights"] == [[1.0, 0.5, 0.25]] + assert batch["teacher_hidden_states"] == [[[1.0, 1.5], [2.0, 2.5], [3.0, 3.5]]] + + +def test_opd_teacher_hidden_states_pad_as_vectors(): + data = [ + { + "input_ids": [1, 2, 3], + "target_tokens": [2, 3, 4], + "teacher_hidden_states": [[1.0, 1.5], [2.0, 2.5], [3.0, 3.5]], + } + ] + packer = SequentialPacker(enable_packing=True, log_stats=False, pad_to_multiple_of=4) + batches = packer.pack(data, max_seq_len=100, request_id="opd-pad") + + assert len(batches) == 1 + assert batches[0]["teacher_hidden_states"] == [[[1.0, 1.5], [2.0, 2.5], [3.0, 3.5], [0.0, 0.0]]] + + def test_packing_exceeds_capacity(simple_data): """Samples overflow one batch -> split into multiple batches.""" packer = SequentialPacker(enable_packing=True, log_stats=False, pad_to_multiple_of=1) @@ -77,7 +151,7 @@ def test_packing_exceeds_capacity(simple_data): def test_packing_disabled(simple_data): - """Packing OFF: one batch per sample.""" + """Packing OFF: one batch per sample, but HF-format datums are still shifted.""" packer = SequentialPacker(enable_packing=False, log_stats=False, pad_to_multiple_of=1) batches = packer.pack(simple_data, max_seq_len=1000, request_id="test-003") @@ -86,6 +160,79 @@ def test_packing_disabled(simple_data): assert batch["batch_id"] == i assert batch["request_id"] == "test-003" assert len(batch["input_ids"]) == 1 + assert len(batch["input_ids"][0]) == len(batch["labels"][0]) == len(batch["position_ids"][0]) + + assert batches[0]["input_ids"] == [[1, 2, 3]] + assert batches[0]["labels"] == [[3, 4, 5]] + assert batches[0]["position_ids"] == [[0, 1, 2]] + assert batches[1]["input_ids"] == [[10]] + assert batches[1]["labels"] == [[30]] + assert batches[1]["position_ids"] == [[0]] + assert batches[2]["input_ids"] == [[100, 200]] + assert batches[2]["labels"] == [[300, 400]] + assert batches[2]["position_ids"] == [[0, 1]] + + +def test_packing_disabled_preserves_shifted_target_tokens(): + """Packing OFF should leave already-shifted xorl_client-format datums unchanged.""" + packer = SequentialPacker(enable_packing=False, log_stats=False, pad_to_multiple_of=1) + datum = { + "model_input": {"input_ids": [11, 22, 33]}, + "loss_fn_inputs": { + "target_tokens": [22, 33, 44], + "logprobs": [0.0, 0.0, 0.0], + "advantages": [1.0, 1.0, 1.0], + }, + } + + batches = packer.pack([datum], max_seq_len=1000, request_id="test-shifted") + + assert len(batches) == 1 + batch = batches[0] + assert batch["input_ids"] == [[11, 22, 33]] + assert batch["labels"] == [[22, 33, 44]] + assert batch["target_tokens"] == [[22, 33, 44]] + assert batch["logprobs"] == [[0.0, 0.0, 0.0]] + assert batch["advantages"] == [[1.0, 1.0, 1.0]] + + +def test_packing_disabled_warns_on_hf_shift(monkeypatch): + """HF labels should warn when shifted in the non-packed path.""" + packer = SequentialPacker(enable_packing=False, log_stats=False, pad_to_multiple_of=1) + warnings = [] + + def _capture_warning(message, *args, **_kwargs): + warnings.append(message % args) + + monkeypatch.setattr("xorl.server.orchestrator.packing.logger.warning", _capture_warning) + datum = { + "input_ids": [1, 2, 3], + "labels": [10, 20, 30], + } + + batches = packer.pack([datum], max_seq_len=1000, request_id="test-shift-warning") + + batch = batches[0] + assert batch["input_ids"] == [[1, 2]] + assert batch["labels"] == [[20, 30]] + assert any("treating it as HF-format data" in warning for warning in warnings) + + +def test_packing_disabled_does_not_overwrite_target_tokens_when_labels_are_present(): + """Preserved target_tokens should not be replaced by labels during non-packed processing.""" + packer = SequentialPacker(enable_packing=False, log_stats=False, pad_to_multiple_of=1) + datum = { + "input_ids": [1, 2, 3], + "labels": [10, 20, 30], + "target_tokens": [101, 102, 103], + } + + batches = packer.pack([datum], max_seq_len=1000, request_id="test-target-preserve") + + batch = batches[0] + assert batch["input_ids"] == [[1, 2, 3]] + assert batch["labels"] == [[10, 20, 30]] + assert batch["target_tokens"] == [[101, 102, 103]] def test_position_ids_and_labels(): diff --git a/tests/server/orchestrator/test_request_processor.py b/tests/server/orchestrator/test_request_processor.py index a8e3f197..fc36d555 100644 --- a/tests/server/orchestrator/test_request_processor.py +++ b/tests/server/orchestrator/test_request_processor.py @@ -13,6 +13,8 @@ - Verify RequestProcessor correctly packs data and formats outputs """ +from unittest.mock import AsyncMock + import pytest import pytest_asyncio @@ -29,8 +31,11 @@ LoadStateData, ModelPassData, OptimStepData, + RegisterSessionData, SaveStateData, + SyncWeightsData, ) +from xorl.server.runner.runner_dispatcher import RunnerDispatcher # ============================================================================ @@ -108,6 +113,206 @@ async def test_forward_backward_operations(processor): assert "loss" in output.outputs[0] +def test_teacher_sort_key_reads_nested_loss_inputs(): + assert RequestProcessor._teacher_sort_key({"loss_fn_inputs": {"teacher_id": 3}}) == 3 + assert RequestProcessor._teacher_sort_key({"loss_fn_inputs": {"teacher_ids": [[2, 2, 2]]}}) == 2 + assert RequestProcessor._teacher_sort_key({"teacher_id": 1, "loss_fn_inputs": {"teacher_id": 4}}) == 1 + + +@pytest.mark.asyncio +async def test_nccl_sync_uses_request_scoped_group_name(): + class CapturingBackend(DummyBackend): + def __init__(self): + super().__init__() + self.group_names = [] + + async def sync_inference_weights(self, *args, **kwargs): + self.group_names.append(kwargs["group_name"]) + return await super().sync_inference_weights(*args, **kwargs) + + backend = CapturingBackend() + exec = RequestProcessor(backend=backend) + await exec.start() + try: + request = OrchestratorRequest( + request_id="sync-req-0001", + request_type=RequestType.ADD, + operation="sync_inference_weights", + payload=SyncWeightsData( + endpoints=[{"host": "127.0.0.1", "port": 30000, "world_size": 1}], + group_name="weight_sync_group", + sync_method="nccl_broadcast", + ), + ) + output = await exec.execute_sync_inference_weights(request) + finally: + await exec.stop() + + assert output.output_type == OutputType.SYNC_INFERENCE_WEIGHTS + assert backend.group_names == ["weight_sync_group_sync_req_0001"] + + +@pytest.mark.asyncio +async def test_model_pass_replay_fields_reach_backend(processor): + """Both routing replay tensors should be forwarded for forward and forward_backward.""" + routed_experts = [[[1, 2], [3, 4]]] + routed_expert_logits = [[[0.1, 0.9], [0.7, 0.3]]] + result = {"total_loss": 1.25, "global_valid_tokens": 3} + processor.backend.forward_backward = AsyncMock(return_value=result) + processor.backend.forward = AsyncMock(return_value=result) + + fb_request = OrchestratorRequest( + request_id="req-r3-fb", + request_type=RequestType.ADD, + operation="forward_backward", + payload=ModelPassData( + data=[{"input_ids": [1, 2, 3], "labels": [2, 3, 4]}], + model_id="session-a", + routed_experts=routed_experts, + routed_expert_logits=routed_expert_logits, + ), + ) + await processor.execute_forward_backward(fb_request) + fb_kwargs = processor.backend.forward_backward.await_args.kwargs + assert fb_kwargs["model_id"] == "session-a" + assert fb_kwargs["routed_experts"] == routed_experts + assert fb_kwargs["routed_expert_logits"] == routed_expert_logits + + fwd_request = OrchestratorRequest( + request_id="req-r3-fwd", + request_type=RequestType.ADD, + operation="forward", + payload=ModelPassData( + data=[{"input_ids": [4, 5, 6], "labels": [5, 6, 7]}], + model_id="session-b", + routed_experts=routed_experts, + routed_expert_logits=routed_expert_logits, + ), + ) + await processor.execute_forward(fwd_request) + fwd_kwargs = processor.backend.forward.await_args.kwargs + assert fwd_kwargs["model_id"] == "session-b" + assert fwd_kwargs["routed_experts"] == routed_experts + assert fwd_kwargs["routed_expert_logits"] == routed_expert_logits + + +def test_runner_dispatcher_forward_compute_preserves_model_id(): + """Forward-only runner execution should switch/use the requested session adapter.""" + + class FakeTrainer: + def __init__(self): + self.forward_kwargs = None + + def forward( + self, + my_batches, + loss_fn, + loss_fn_params, + *, + model_id="default", + routed_experts=None, + routed_expert_logits=None, + ): + self.forward_kwargs = { + "my_batches": my_batches, + "loss_fn": loss_fn, + "loss_fn_params": loss_fn_params, + "model_id": model_id, + "routed_experts": routed_experts, + "routed_expert_logits": routed_expert_logits, + } + return {"success": True, "model_id": model_id} + + dispatcher = object.__new__(RunnerDispatcher) + dispatcher.trainer = FakeTrainer() + routed_experts = [[[1, 2]]] + routed_expert_logits = [[[0.25, 0.75]]] + + result = RunnerDispatcher._execute_compute( + dispatcher, + [{"input_ids": [1, 2], "labels": [2, 3]}], + "causallm_loss", + {"return_per_token": False}, + routed_experts, + with_backward=False, + model_id="session-a", + routed_expert_logits=routed_expert_logits, + ) + + assert result["model_id"] == "session-a" + assert dispatcher.trainer.forward_kwargs["model_id"] == "session-a" + assert dispatcher.trainer.forward_kwargs["routed_experts"] == routed_experts + assert dispatcher.trainer.forward_kwargs["routed_expert_logits"] == routed_expert_logits + + +@pytest.mark.asyncio +async def test_runner_dispatcher_forward_rank0_scatter_preserves_model_id(): + """The rank-0 forward handler must not drop model_id before compute execution.""" + + class FakeCoordinator: + def auto_load_if_evicted(self, model_id): + captured["auto_load_model_id"] = model_id + return False, None + + captured = {} + dispatcher = object.__new__(RunnerDispatcher) + dispatcher._adapter_coordinator = FakeCoordinator() + + routed_experts = [[[1, 2]]] + routed_expert_logits = [[[0.25, 0.75]]] + + def select_batches(batches, routed_experts=None, routed_expert_logits=None): + return batches, routed_experts, routed_expert_logits + + def execute_and_gather( + my_batches, + loss_fn, + loss_fn_params, + routed_experts, + cp_enabled, + parallel_state, + *, + with_backward, + model_id, + is_rank0, + routed_expert_logits=None, + ): + captured.update( + { + "model_id": model_id, + "with_backward": with_backward, + "is_rank0": is_rank0, + "routed_experts": routed_experts, + "routed_expert_logits": routed_expert_logits, + } + ) + return {"success": True, "model_id": model_id} + + dispatcher._select_and_prepare_batches = select_batches + dispatcher._execute_and_gather = execute_and_gather + + result = await RunnerDispatcher._handle_compute_rank0_scatter( + dispatcher, + { + "payload": ModelPassData( + batches=[{"input_ids": [1, 2], "labels": [2, 3]}], + model_id="session-a", + routed_experts=routed_experts, + routed_expert_logits=routed_expert_logits, + ) + }, + with_backward=False, + ) + + assert result["model_id"] == "session-a" + assert captured["model_id"] == "session-a" + assert captured["auto_load_model_id"] == "session-a" + assert captured["with_backward"] is False + assert captured["is_rank0"] is True + assert captured["routed_experts"] == routed_experts + assert captured["routed_expert_logits"] == routed_expert_logits + + @pytest.mark.asyncio async def test_optim_and_checkpoint_operations(processor): """Test optim_step, save_state, load_state, sleep, and wake_up.""" @@ -166,6 +371,81 @@ async def test_optim_and_checkpoint_operations(processor): assert output.output_type == OutputType.WAKE_UP +@pytest.mark.asyncio +async def test_register_session_operation_reaches_backend(processor): + """register_session should flow through the processor to the backend.""" + session_spec = { + "base_model": "Qwen/Qwen3-8B", + "lora_config": {"lora_rank": 4, "lora_alpha": 8}, + "optimizer_config": {"type": "adamw", "learning_rate": 1e-4}, + } + request = OrchestratorRequest( + request_id="req-register-session", + request_type=RequestType.ADD, + operation="register_session", + payload=RegisterSessionData(model_id="session-a", session_spec=session_spec, materialize=True), + ) + + output = await processor.execute_register_session(request) + + assert output.output_type == OutputType.REGISTER_SESSION + result = output.outputs["result"] + assert result["registered"] is True + assert result["model_id"] == "session-a" + assert result["session_spec"] == session_spec + assert result["materialize"] is True + + +@pytest.mark.asyncio +async def test_runner_dispatcher_register_session_handler_materializes_adapter(): + """Remote register_session should be a real runner operation, not an unknown command.""" + + class FakeCoordinator: + def __init__(self): + self.command_dict = None + + async def handle_register_session(self, command_dict): + self.command_dict = command_dict + payload = command_dict["payload"] + lr = payload.session_spec["optimizer_config"]["learning_rate"] + return { + "registered": True, + "model_id": payload.model_id, + "lr": lr, + "session_spec": payload.session_spec, + "materialize": payload.materialize, + } + + dispatcher = object.__new__(RunnerDispatcher) + dispatcher.rank = 0 + dispatcher._adapter_coordinator = FakeCoordinator() + session_spec = { + "optimizer_config": {"learning_rate": 2e-4}, + "lora_config": {"lora_rank": 4, "lora_alpha": 8}, + } + + result = await RunnerDispatcher._handle_register_session( + dispatcher, + { + "payload": RegisterSessionData( + model_id="session-a", + session_spec=session_spec, + materialize=True, + ) + }, + ) + + assert RunnerDispatcher._COMMAND_HANDLERS["register_session"] == "_handle_register_session" + assert result["registered"] is True + assert result["model_id"] == "session-a" + assert result["lr"] == pytest.approx(2e-4) + assert result["session_spec"] == session_spec + assert result["materialize"] is True + forwarded_payload = dispatcher._adapter_coordinator.command_dict["payload"] + assert forwarded_payload.session_spec["optimizer_config"]["learning_rate"] == pytest.approx(2e-4) + assert forwarded_payload.materialize is True + + @pytest.mark.asyncio async def test_statistics_tracking(processor): """Test that statistics track operations correctly.""" diff --git a/tests/server/runner/test_adapter_coordinator.py b/tests/server/runner/test_adapter_coordinator.py new file mode 100644 index 00000000..255d534c --- /dev/null +++ b/tests/server/runner/test_adapter_coordinator.py @@ -0,0 +1,645 @@ +"""Tests for multi-rank adapter load coordination.""" + +import asyncio +import importlib.util +import json +import time +from pathlib import Path +from unittest.mock import Mock + +import pytest +import torch +from safetensors.torch import save_file as save_safetensors_file + +from xorl.lora.utils import LoraTensorShardSpec +from xorl.server.protocol.operations import AdapterStateData, RegisterAdapterData, RegisterSessionData + + +_MODULE_PATH = ( + Path(__file__).resolve().parents[3] / "src" / "xorl" / "server" / "runner" / "adapters" / "adapter_coordinator.py" +) +_SPEC = importlib.util.spec_from_file_location("xorl_test_adapter_coordinator", _MODULE_PATH) +assert _SPEC is not None and _SPEC.loader is not None +_MODULE = importlib.util.module_from_spec(_SPEC) +_SPEC.loader.exec_module(_MODULE) +AdapterCoordinator = _MODULE.AdapterCoordinator + + +pytestmark = [pytest.mark.cpu, pytest.mark.server] + + +class _FakeAdapterState: + def __init__(self, lr: float = 1e-5): + self.global_step = 0 + self.global_forward_backward_step = 0 + self.lr = lr + self.lora_params = {} + self.last_access_time = time.time() + + +class _FakeAdapterManager: + def __init__(self, checkpoint_dir: Path): + self.checkpoint_dir = str(checkpoint_dir) + self.adapters = {} + self.current_adapter_id = None + self.max_adapters = 8 + self.model = None + + @staticmethod + def _canonical_lora_param_name(name: str) -> str: + if name.endswith(".weight"): + return name[: -len(".weight")] + return name + + def get_adapter_state(self, model_id: str): + return self.adapters[model_id] + + def has_adapter(self, model_id: str) -> bool: + return model_id in self.adapters + + def remove_adapter(self, model_id: str) -> None: + self.adapters.pop(model_id, None) + + def list_adapters(self): + return list(self.adapters.keys()) + + def set_lr(self, model_id: str, lr: float) -> None: + self.adapters[model_id].lr = lr + + +class _FakeTrainer: + def __init__(self, checkpoint_dir: Path, *, adapter_state_load_mode: str = "all_ranks"): + self.adapter_manager = _FakeAdapterManager(checkpoint_dir) + self.register_calls = [] + self.load_calls = [] + self.save_calls = [] + self.lora_config = {"adapter_state_load_mode": adapter_state_load_mode} + self.train_config = {"pipeline_parallel_size": 1} + self.lora_session_specs = {} + self.fail_load_with = None + + @staticmethod + def _session_spec(lr: float = 1e-5): + return { + "base_model": "Qwen/Qwen3-8B", + "is_lora": True, + "lora_config": {"lora_rank": 4, "lora_alpha": 8}, + "optimizer_config": { + "type": "adamw", + "learning_rate": lr, + "weight_decay": 0.01, + "optimizer_dtype": "bf16", + "betas": [0.9, 0.95], + "eps": 1e-8, + "optimizer_kwargs": {}, + }, + } + + def get_lora_session_spec(self, model_id: str): + return self.lora_session_specs[model_id] + + def register_session( + self, model_id: str, session_spec: dict, materialize: bool = False, initialize_fresh: bool = True + ): + self.lora_session_specs[model_id] = session_spec + if materialize: + self.register_lora_adapter(model_id, session_spec["optimizer_config"]["learning_rate"]) + return {"registered": True, "model_id": model_id, "materialized": materialize} + + def register_lora_adapter(self, model_id: str, lr: float): + self.lora_session_specs.setdefault(model_id, self._session_spec(lr)) + self.register_calls.append((model_id, lr)) + self.adapter_manager.adapters[model_id] = _FakeAdapterState(lr=lr) + + def register_adapter(self, model_id: str, lr: float): + self.register_lora_adapter(model_id, lr) + return {"registered": True, "model_id": model_id, "lr": lr} + + def load_adapter_state(self, model_id: str, path: str, load_optimizer: bool = True, lr: float | None = None): + if self.fail_load_with is not None: + raise self.fail_load_with + self.load_calls.append( + { + "model_id": model_id, + "path": path, + "load_optimizer": load_optimizer, + "lr": lr, + } + ) + state = self.adapter_manager.adapters.setdefault(model_id, _FakeAdapterState()) + if lr is not None: + state.lr = lr + state.global_step = 7 + return {"success": True, "model_id": model_id, "step": state.global_step} + + def save_adapter_state(self, model_id: str, path: str, save_optimizer: bool = True): + self.save_calls.append( + { + "model_id": model_id, + "path": path, + "save_optimizer": save_optimizer, + } + ) + return {"success": True, "model_id": model_id, "path": path} + + +def test_auto_load_if_evicted_loads_adapter_state_on_non_rank0(tmp_path): + checkpoint_dir = tmp_path / "adapters" + evicted_path = checkpoint_dir / "evicted" / "policy-a" + evicted_path.mkdir(parents=True) + + trainer = _FakeTrainer(checkpoint_dir) + trainer.register_session("policy-a", trainer._session_spec(1e-5), materialize=False) + coordinator = AdapterCoordinator(trainer=trainer, rank=1, world_size=2, cpu_group=None) + coordinator.broadcast_adapter_state = Mock() + coordinator.broadcast_adapter_optimizer_state = Mock() + + was_auto_loaded, checkpoint_path = coordinator.auto_load_if_evicted("policy-a") + + assert was_auto_loaded is True + assert checkpoint_path == str(evicted_path) + assert trainer.register_calls == [("policy-a", 1e-5)] + assert trainer.load_calls == [ + { + "model_id": "policy-a", + "path": str(evicted_path), + "load_optimizer": True, + "lr": None, + } + ] + coordinator.broadcast_adapter_state.assert_not_called() + coordinator.broadcast_adapter_optimizer_state.assert_not_called() + + +def test_auto_load_if_evicted_broadcasts_fresh_adapter_without_checkpoint(tmp_path): + checkpoint_dir = tmp_path / "adapters" + checkpoint_dir.mkdir() + + trainer = _FakeTrainer(checkpoint_dir) + trainer.register_session("policy-new", trainer._session_spec(1e-5), materialize=False) + coordinator = AdapterCoordinator(trainer=trainer, rank=1, world_size=2, cpu_group=None) + coordinator.broadcast_adapter_state = Mock() + coordinator.broadcast_adapter_optimizer_state = Mock() + + was_auto_loaded, checkpoint_path = coordinator.auto_load_if_evicted("policy-new") + + assert was_auto_loaded is True + assert checkpoint_path is None + assert trainer.register_calls == [("policy-new", 1e-5)] + assert trainer.load_calls == [] + coordinator.broadcast_adapter_state.assert_called_once_with("policy-new", 1e-5) + coordinator.broadcast_adapter_optimizer_state.assert_not_called() + + +def test_auto_load_if_evicted_syncs_fresh_materialization_failure_before_broadcast(tmp_path): + checkpoint_dir = tmp_path / "adapters" + checkpoint_dir.mkdir() + + trainer = _FakeTrainer(checkpoint_dir) + trainer.register_session("policy-fail", trainer._session_spec(1e-5), materialize=False) + + def _fail_register(model_id, lr): + raise RuntimeError("capacity full") + + trainer.register_lora_adapter = _fail_register + coordinator = AdapterCoordinator(trainer=trainer, rank=1, world_size=2, cpu_group=None) + coordinator.broadcast_adapter_state = Mock() + coordinator._sync_collective_error = Mock(return_value="rank 1: capacity full") + + with pytest.raises(RuntimeError, match="rank 1: capacity full"): + coordinator.auto_load_if_evicted("policy-fail") + + coordinator._sync_collective_error.assert_called_once() + coordinator.broadcast_adapter_state.assert_not_called() + + +def test_auto_load_if_evicted_rejects_missing_checkpoint_when_fresh_materialization_disabled(tmp_path): + checkpoint_dir = tmp_path / "adapters" + checkpoint_dir.mkdir() + + trainer = _FakeTrainer(checkpoint_dir) + trainer.register_session("policy-missing", trainer._session_spec(2e-5), materialize=False) + coordinator = AdapterCoordinator(trainer=trainer, rank=0, world_size=1, cpu_group=None) + coordinator.broadcast_adapter_state = Mock() + + with pytest.raises(FileNotFoundError, match="Refusing to recreate fresh state"): + coordinator.auto_load_if_evicted("policy-missing", allow_fresh_materialization=False) + + assert trainer.register_calls == [] + coordinator.broadcast_adapter_state.assert_not_called() + + +def test_auto_load_if_evicted_rolls_back_fresh_adapter_when_restore_fails(tmp_path): + checkpoint_dir = tmp_path / "adapters" + evicted_path = checkpoint_dir / "evicted" / "policy-bad" + evicted_path.mkdir(parents=True) + + trainer = _FakeTrainer(checkpoint_dir) + trainer.fail_load_with = RuntimeError("corrupt checkpoint") + trainer.register_session("policy-bad", trainer._session_spec(2e-5), materialize=False) + coordinator = AdapterCoordinator(trainer=trainer, rank=0, world_size=1, cpu_group=None) + coordinator.broadcast_adapter_state = Mock() + + with pytest.raises(RuntimeError, match="Failed to auto-load adapter 'policy-bad'"): + coordinator.auto_load_if_evicted("policy-bad") + + assert trainer.register_calls == [("policy-bad", 2e-5)] + assert trainer.load_calls == [] + assert not trainer.adapter_manager.has_adapter("policy-bad") + coordinator.broadcast_adapter_state.assert_not_called() + + +def test_handle_load_adapter_state_loads_optimizer_on_non_rank0(tmp_path): + checkpoint_dir = tmp_path / "adapters" + checkpoint_dir.mkdir() + adapter_path = tmp_path / "checkpoint" + adapter_path.mkdir() + + trainer = _FakeTrainer(checkpoint_dir) + trainer.register_session("policy-b", trainer._session_spec(3e-5), materialize=False) + coordinator = AdapterCoordinator(trainer=trainer, rank=1, world_size=2, cpu_group=None) + coordinator.broadcast_adapter_state = Mock() + coordinator.broadcast_adapter_optimizer_state = Mock() + + payload = AdapterStateData( + model_id="policy-b", + path=str(adapter_path), + load_optimizer=True, + lr=3e-5, + ) + + result = asyncio.run(coordinator.handle_load_adapter_state({"payload": payload})) + + assert result == {"success": True, "model_id": "policy-b"} + assert trainer.register_calls == [("policy-b", 3e-5)] + assert trainer.load_calls == [ + { + "model_id": "policy-b", + "path": str(adapter_path), + "load_optimizer": True, + "lr": 3e-5, + } + ] + coordinator.broadcast_adapter_state.assert_not_called() + coordinator.broadcast_adapter_optimizer_state.assert_not_called() + + +def test_auto_load_if_evicted_uses_rank0_broadcast_mode_on_non_rank0(tmp_path): + checkpoint_dir = tmp_path / "adapters" + evicted_path = checkpoint_dir / "evicted" / "policy-c" + evicted_path.mkdir(parents=True) + + trainer = _FakeTrainer(checkpoint_dir, adapter_state_load_mode="rank0_broadcast") + trainer.register_session("policy-c", trainer._session_spec(1e-5), materialize=False) + coordinator = AdapterCoordinator(trainer=trainer, rank=1, world_size=2, cpu_group=None) + coordinator.broadcast_adapter_state = Mock() + coordinator.broadcast_adapter_optimizer_state = Mock() + + was_auto_loaded, checkpoint_path = coordinator.auto_load_if_evicted("policy-c") + + assert was_auto_loaded is True + assert checkpoint_path == str(evicted_path) + assert trainer.register_calls == [("policy-c", 1e-5)] + assert trainer.load_calls == [] + coordinator.broadcast_adapter_state.assert_called_once_with("policy-c", 1e-5) + coordinator.broadcast_adapter_optimizer_state.assert_called_once_with("policy-c") + + +def test_handle_load_adapter_state_uses_rank0_broadcast_mode_on_non_rank0(tmp_path): + checkpoint_dir = tmp_path / "adapters" + checkpoint_dir.mkdir() + adapter_path = tmp_path / "checkpoint" + adapter_path.mkdir() + + trainer = _FakeTrainer(checkpoint_dir, adapter_state_load_mode="rank0_broadcast") + trainer.register_session("policy-d", trainer._session_spec(2e-5), materialize=False) + coordinator = AdapterCoordinator(trainer=trainer, rank=1, world_size=2, cpu_group=None) + coordinator.broadcast_adapter_state = Mock() + coordinator.broadcast_adapter_optimizer_state = Mock() + + payload = AdapterStateData( + model_id="policy-d", + path=str(adapter_path), + load_optimizer=True, + lr=2e-5, + ) + + result = asyncio.run(coordinator.handle_load_adapter_state({"payload": payload})) + + assert result == {"success": True, "model_id": "policy-d"} + assert trainer.register_calls == [("policy-d", 2e-5)] + assert trainer.load_calls == [] + coordinator.broadcast_adapter_state.assert_called_once_with("policy-d", 2e-5) + coordinator.broadcast_adapter_optimizer_state.assert_called_once_with("policy-d") + + +def test_rank0_broadcast_ep_sharded_restore_slices_full_checkpoint_tensor(monkeypatch, tmp_path): + checkpoint_dir = tmp_path / "adapters" + checkpoint_dir.mkdir() + adapter_path = tmp_path / "checkpoint" + adapter_path.mkdir() + + param_name = "model.layers.0.mlp.experts.down_proj_lora_A" + full_tensor = torch.arange(4 * 5 * 2, dtype=torch.float32).reshape(4, 5, 2) + save_safetensors_file( + { + f"base_model.model.model.layers.0.mlp.experts.{expert_idx}.down_proj.lora_A.weight": full_tensor[expert_idx] + .transpose(0, 1) + .contiguous() + for expert_idx in range(4) + }, + str(adapter_path / "adapter_model.safetensors"), + ) + (adapter_path / "metadata.json").write_text( + json.dumps({"global_step": 11, "global_forward_backward_step": 13, "lr": 3e-5}), + encoding="utf-8", + ) + + trainer = _FakeTrainer(checkpoint_dir, adapter_state_load_mode="rank0_broadcast") + trainer.register_session("policy-ep", trainer._session_spec(3e-5), materialize=False) + trainer.register_lora_adapter("policy-ep", 3e-5) + trainer.adapter_manager.model = object() + state = trainer.adapter_manager.get_adapter_state("policy-ep") + state.lora_params = {param_name: torch.nn.Parameter(torch.empty(2, 5, 2))} + + monkeypatch.setattr( + _MODULE, + "get_lora_tensor_shard_specs", + lambda model, names=None: {param_name: LoraTensorShardSpec(dim=0, index=1, size=2)}, + ) + monkeypatch.setattr(_MODULE.dist, "broadcast_object_list", lambda payload, src=0, group=None: None) + + coordinator = AdapterCoordinator(trainer=trainer, rank=0, world_size=2, cpu_group=None) + coordinator.broadcast_adapter_state = Mock() + coordinator.broadcast_adapter_optimizer_state = Mock() + + result = coordinator._restore_adapter_state( + model_id="policy-ep", + path=str(adapter_path), + load_optimizer=False, + lr=None, + default_lr=3e-5, + ) + + assert result["success"] is True + assert result["step"] == 11 + assert torch.equal(state.lora_params[param_name].detach(), full_tensor[2:4]) + assert state.global_forward_backward_step == 13 + assert state.lr == 3e-5 + assert trainer.load_calls == [] + coordinator.broadcast_adapter_state.assert_not_called() + coordinator.broadcast_adapter_optimizer_state.assert_not_called() + + +def test_rank0_broadcast_ep_sharded_restore_rejects_session_spec_mismatch(monkeypatch, tmp_path): + checkpoint_dir = tmp_path / "adapters" + checkpoint_dir.mkdir() + adapter_path = tmp_path / "checkpoint" + adapter_path.mkdir() + + param_name = "model.layers.0.mlp.experts.down_proj_lora_A" + full_tensor = torch.arange(4 * 5 * 2, dtype=torch.float32).reshape(4, 5, 2) + save_safetensors_file( + { + f"base_model.model.model.layers.0.mlp.experts.{expert_idx}.down_proj.lora_A.weight": full_tensor[expert_idx] + .transpose(0, 1) + .contiguous() + for expert_idx in range(4) + }, + str(adapter_path / "adapter_model.safetensors"), + ) + (adapter_path / "metadata.json").write_text( + json.dumps({"global_step": 11, "global_forward_backward_step": 13, "lr": 3e-5}), + encoding="utf-8", + ) + checkpoint_session_spec = _FakeTrainer._session_spec(3e-5) + checkpoint_session_spec["lora_config"]["lora_alpha"] = 16 + (adapter_path / "session_spec.json").write_text(json.dumps(checkpoint_session_spec), encoding="utf-8") + + trainer = _FakeTrainer(checkpoint_dir, adapter_state_load_mode="rank0_broadcast") + trainer.register_session("policy-ep", trainer._session_spec(3e-5), materialize=False) + trainer.register_lora_adapter("policy-ep", 3e-5) + trainer.adapter_manager.model = object() + state = trainer.adapter_manager.get_adapter_state("policy-ep") + original_tensor = torch.zeros(2, 5, 2) + state.lora_params = {param_name: torch.nn.Parameter(original_tensor.clone())} + + monkeypatch.setattr( + _MODULE, + "get_lora_tensor_shard_specs", + lambda model, names=None: {param_name: LoraTensorShardSpec(dim=0, index=1, size=2)}, + ) + monkeypatch.setattr(_MODULE.dist, "broadcast_object_list", lambda payload, src=0, group=None: None) + + coordinator = AdapterCoordinator(trainer=trainer, rank=0, world_size=2, cpu_group=None) + + with pytest.raises(ValueError, match="Checkpoint session spec does not match"): + coordinator._restore_ep_sharded_rank0_broadcast_adapter_state( + model_id="policy-ep", + path=str(adapter_path), + load_optimizer=True, + lr=None, + ) + + assert torch.equal(state.lora_params[param_name].detach(), original_tensor) + + +def test_handle_load_adapter_state_rolls_back_new_adapter_on_cross_rank_restore_error(tmp_path): + checkpoint_dir = tmp_path / "adapters" + checkpoint_dir.mkdir() + adapter_path = tmp_path / "checkpoint" + adapter_path.mkdir() + + trainer = _FakeTrainer(checkpoint_dir, adapter_state_load_mode="rank0_broadcast") + trainer.register_session("policy-sync-fail", trainer._session_spec(2e-5), materialize=False) + coordinator = AdapterCoordinator(trainer=trainer, rank=1, world_size=2, cpu_group=None) + coordinator.broadcast_adapter_state = Mock() + coordinator.broadcast_adapter_optimizer_state = Mock() + coordinator._sync_collective_error = Mock(side_effect=[None, None, "rank 0: Adapter state restore failed"]) + + payload = AdapterStateData( + model_id="policy-sync-fail", + path=str(adapter_path), + load_optimizer=True, + lr=2e-5, + ) + + result = asyncio.run(coordinator.handle_load_adapter_state({"payload": payload})) + + assert result == { + "success": False, + "error": "Adapter state load failed: rank 0: Adapter state restore failed", + } + assert trainer.register_calls == [("policy-sync-fail", 2e-5)] + assert not trainer.adapter_manager.has_adapter("policy-sync-fail") + coordinator.broadcast_adapter_state.assert_not_called() + coordinator.broadcast_adapter_optimizer_state.assert_not_called() + + +def test_handle_load_adapter_state_rejects_pipeline_parallel_multi_adapter_lora(tmp_path): + checkpoint_dir = tmp_path / "adapters" + checkpoint_dir.mkdir() + adapter_path = tmp_path / "checkpoint" + adapter_path.mkdir() + + trainer = _FakeTrainer(checkpoint_dir, adapter_state_load_mode="rank0_broadcast") + trainer.train_config["pipeline_parallel_size"] = 2 + trainer.register_session("policy-pp", trainer._session_spec(2e-5), materialize=False) + coordinator = AdapterCoordinator(trainer=trainer, rank=1, world_size=2, cpu_group=None) + + payload = AdapterStateData( + model_id="policy-pp", + path=str(adapter_path), + load_optimizer=True, + lr=2e-5, + ) + + result = asyncio.run(coordinator.handle_load_adapter_state({"payload": payload})) + + assert result == { + "success": False, + "error": ( + "Adapter state load failed: pipeline_parallel_size > 1 is not supported with multi-adapter LoRA " + "server training. Adapter coordination currently assumes identical local LoRA layouts on every rank." + ), + } + assert "policy-pp" in trainer.lora_session_specs + assert trainer.load_calls == [] + + +def test_handle_load_adapter_state_rolls_back_auto_registered_session_on_failure(tmp_path): + checkpoint_dir = tmp_path / "adapters" + checkpoint_dir.mkdir() + adapter_path = tmp_path / "checkpoint" + adapter_path.mkdir() + (adapter_path / "session_spec.json").write_text(json.dumps(_FakeTrainer._session_spec(4e-5)), encoding="utf-8") + + trainer = _FakeTrainer(checkpoint_dir) + trainer.fail_load_with = RuntimeError("corrupt checkpoint") + coordinator = AdapterCoordinator(trainer=trainer, rank=0, world_size=1, cpu_group=None) + coordinator.broadcast_adapter_state = Mock() + + payload = AdapterStateData( + model_id="policy-checkpoint-only", + path=str(adapter_path), + load_optimizer=True, + ) + + result = asyncio.run(coordinator.handle_load_adapter_state({"payload": payload})) + + assert result["success"] is False + assert result["error"] == ( + "Adapter state load failed: Adapter state restore failed for model_id=policy-checkpoint-only: corrupt checkpoint" + ) + assert "policy-checkpoint-only" not in trainer.lora_session_specs + assert not trainer.adapter_manager.has_adapter("policy-checkpoint-only") + coordinator.broadcast_adapter_state.assert_not_called() + + +def test_handle_register_adapter_broadcasts_fresh_adapter_state(tmp_path): + checkpoint_dir = tmp_path / "adapters" + checkpoint_dir.mkdir() + + trainer = _FakeTrainer(checkpoint_dir) + coordinator = AdapterCoordinator(trainer=trainer, rank=1, world_size=2, cpu_group=None) + coordinator.broadcast_adapter_state = Mock() + + result = asyncio.run(coordinator.handle_register_adapter({"payload": RegisterAdapterData("policy-e", 4e-5)})) + + assert result == {} + assert trainer.register_calls == [("policy-e", 4e-5)] + coordinator.broadcast_adapter_state.assert_called_once_with("policy-e", 4e-5) + + +def test_handle_register_session_materializes_and_broadcasts(tmp_path): + checkpoint_dir = tmp_path / "adapters" + checkpoint_dir.mkdir() + + trainer = _FakeTrainer(checkpoint_dir) + coordinator = AdapterCoordinator(trainer=trainer, rank=1, world_size=2, cpu_group=None) + coordinator.broadcast_adapter_state = Mock() + + session_spec = trainer._session_spec(5e-5) + result = asyncio.run( + coordinator.handle_register_session( + {"payload": RegisterSessionData(model_id="policy-f", session_spec=session_spec, materialize=True)} + ) + ) + + assert result == {} + assert trainer.lora_session_specs["policy-f"] == session_spec + assert trainer.register_calls == [("policy-f", 5e-5)] + coordinator.broadcast_adapter_state.assert_called_once_with("policy-f", 5e-5) + + +def test_handle_register_session_rolls_back_new_state_on_cross_rank_failure(tmp_path): + checkpoint_dir = tmp_path / "adapters" + checkpoint_dir.mkdir() + + trainer = _FakeTrainer(checkpoint_dir) + coordinator = AdapterCoordinator(trainer=trainer, rank=1, world_size=2, cpu_group=None) + coordinator.broadcast_adapter_state = Mock() + coordinator._sync_collective_error = Mock(return_value="rank 0: pending gradients") + + session_spec = trainer._session_spec(5e-5) + with pytest.raises(RuntimeError, match="Session registration failed: rank 0: pending gradients"): + asyncio.run( + coordinator.handle_register_session( + {"payload": RegisterSessionData(model_id="policy-fail", session_spec=session_spec, materialize=True)} + ) + ) + + assert "policy-fail" not in trainer.lora_session_specs + assert not trainer.adapter_manager.has_adapter("policy-fail") + coordinator.broadcast_adapter_state.assert_not_called() + + +def test_handle_register_session_raises_when_worker_registration_fails(tmp_path): + checkpoint_dir = tmp_path / "adapters" + checkpoint_dir.mkdir() + + trainer = _FakeTrainer(checkpoint_dir) + + def _fail_register(*args, **kwargs): + raise ValueError("boom") + + trainer.register_session = _fail_register + coordinator = AdapterCoordinator(trainer=trainer, rank=1, world_size=2, cpu_group=None) + + with pytest.raises(RuntimeError, match="Session registration failed: boom"): + asyncio.run( + coordinator.handle_register_session( + { + "payload": RegisterSessionData( + model_id="policy-g", + session_spec=trainer._session_spec(3e-5), + materialize=False, + ) + } + ) + ) + + +def test_handle_save_adapter_state_requires_evicted_checkpoint(tmp_path): + checkpoint_dir = tmp_path / "adapters" + checkpoint_dir.mkdir() + + trainer = _FakeTrainer(checkpoint_dir) + trainer.register_session("policy-save", trainer._session_spec(4e-5), materialize=False) + coordinator = AdapterCoordinator(trainer=trainer, rank=0, world_size=1, cpu_group=None) + + with pytest.raises(RuntimeError, match="Adapter state save failed: .*Refusing to recreate fresh state"): + asyncio.run( + coordinator.handle_save_adapter_state( + { + "payload": AdapterStateData( + model_id="policy-save", + path=str(tmp_path / "save-target"), + save_optimizer=True, + ) + } + ) + ) + + assert trainer.register_calls == [] + assert trainer.save_calls == [] diff --git a/tests/server/runner/test_adapter_manager.py b/tests/server/runner/test_adapter_manager.py new file mode 100644 index 00000000..c234ff0a --- /dev/null +++ b/tests/server/runner/test_adapter_manager.py @@ -0,0 +1,613 @@ +"""Tests for adapter-manager optimizer integration.""" + +import asyncio +import json +from pathlib import Path + +import pytest +import torch +import torch.nn as nn +from safetensors.torch import load_file as safetensors_load_file +from safetensors.torch import save_file as safetensors_save_file + +from xorl.optim import SignSGD +from xorl.server.protocol.operations import AdapterStateData +from xorl.server.runner.adapters.adapter_coordinator import AdapterCoordinator +from xorl.server.runner.adapters.manager import LoRAAdapterManager +from xorl.server.session_spec import normalize_session_spec + + +pytestmark = [pytest.mark.cpu, pytest.mark.server] + + +class _DummyLoRALayer(nn.Module): + def __init__(self, *, max_rank: int = 4) -> None: + super().__init__() + self.lora_A = nn.Parameter(torch.randn(max_rank, 8)) + self.lora_B = nn.Parameter(torch.zeros(8, max_rank)) + self.active_r = max_rank + self.active_lora_alpha = 16 + + def set_runtime_lora_config(self, lora_rank: int, lora_alpha: int) -> None: + self.active_r = lora_rank + self.active_lora_alpha = lora_alpha + + +class _DummyLoRAModel(nn.Module): + def __init__(self, *, max_rank: int = 4) -> None: + super().__init__() + self.model = nn.Module() + self.model.layers = nn.ModuleList([nn.Module()]) + self.model.layers[0].self_attn = nn.Module() + self.model.layers[0].self_attn.o_proj = _DummyLoRALayer(max_rank=max_rank) + + +def _build_manager(tmp_path: Path, **kwargs) -> LoRAAdapterManager: + max_rank = kwargs.pop("max_rank", 4) + return LoRAAdapterManager( + _DummyLoRAModel(max_rank=max_rank), + device=torch.device("cpu"), + checkpoint_dir=str(tmp_path / "adapters"), + auto_save_on_eviction=False, + **kwargs, + ) + + +class _CoordinatorTrainer: + def __init__(self, adapter_manager: LoRAAdapterManager) -> None: + self.adapter_manager = adapter_manager + self.lora_session_specs = {} + + def register_session(self, model_id: str, session_spec: dict, materialize: bool = False, **kwargs): + self.lora_session_specs[model_id] = session_spec + if materialize and not self.adapter_manager.has_adapter(model_id): + self.adapter_manager.register_adapter( + model_id=model_id, + session_spec=session_spec, + initialize_fresh=kwargs.get("initialize_fresh", True), + ) + return {"registered": True, "model_id": model_id} + + def get_lora_session_spec(self, model_id: str) -> dict: + if model_id in self.lora_session_specs: + return self.lora_session_specs[model_id] + return self.adapter_manager.get_adapter_session_spec(model_id) + + def register_lora_adapter(self, model_id: str, lr=None): + session_spec = dict(self.get_lora_session_spec(model_id)) + if lr is not None: + session_spec.setdefault("optimizer_config", {})["learning_rate"] = lr + self.adapter_manager.register_adapter(model_id=model_id, session_spec=session_spec, initialize_fresh=True) + return {"registered": True, "model_id": model_id} + + def load_adapter_state(self, model_id: str, path: str, load_optimizer: bool = True, lr=None): + return self.adapter_manager.load_adapter_state( + model_id=model_id, + path=path, + load_optimizer=load_optimizer, + lr=lr, + ) + + +def _session_spec(*, rank: int, alpha: int, optimizer_type: str, lr: float, weight_decay: float = 0.0) -> dict: + return { + "base_model": "Qwen/Qwen3-8B", + "is_lora": True, + "lora_config": { + "lora_rank": rank, + "lora_alpha": alpha, + }, + "optimizer_config": { + "type": optimizer_type, + "learning_rate": lr, + "weight_decay": weight_decay, + "optimizer_dtype": "bf16", + "betas": None if optimizer_type in {"sgd", "signsgd"} else [0.9, 0.95], + "eps": None if optimizer_type in {"sgd", "signsgd"} else 1e-8, + "optimizer_kwargs": {}, + }, + } + + +def test_register_adapter_uses_shared_optimizer_factory_and_checkpoint_dir(tmp_path): + manager = _build_manager(tmp_path, optimizer_type="signsgd", weight_decay=0.25) + + manager.register_adapter("policy-a", lr=0.1, initialize_fresh=True) + state = manager.get_adapter_state("policy-a") + + assert isinstance(state.optimizer, SignSGD) + + save_result = manager.save_adapter_state("policy-a") + save_path = Path(save_result["path"]) + metadata = json.loads((save_path / "metadata.json").read_text(encoding="utf-8")) + + assert save_path == tmp_path / "adapters" / "policy-a" + assert metadata["optimizer"]["type"] == "signsgd" + assert metadata["optimizer"]["weight_decay"] == pytest.approx(0.25) + assert metadata["optimizer"]["betas"] == [0.9, 0.95] + assert metadata["optimizer"]["eps"] == pytest.approx(1e-8) + + +def test_save_adapter_state_preserves_lora_weight_dtype(tmp_path): + manager = _build_manager(tmp_path, optimizer_type="adamw") + manager.register_adapter("policy-fp32", lr=0.1, initialize_fresh=True) + + save_path = Path(manager.save_adapter_state("policy-fp32")["path"]) + weights = safetensors_load_file(str(save_path / "adapter_model.safetensors")) + + assert weights["base_model.model.model.layers.0.self_attn.o_proj.lora_A"].dtype == torch.float32 + assert weights["base_model.model.model.layers.0.self_attn.o_proj.lora_B"].dtype == torch.float32 + + +def test_load_adapter_state_uses_checkpoint_optimizer_contract_for_fresh_session(tmp_path): + source_manager = _build_manager(tmp_path, optimizer_type="signsgd") + source_manager.register_adapter("policy-b", lr=0.1, initialize_fresh=True) + checkpoint_path = source_manager.save_adapter_state("policy-b")["path"] + + target_manager = _build_manager(tmp_path, optimizer_type="adamw") + result = target_manager.load_adapter_state("policy-b", checkpoint_path, load_optimizer=True) + + assert result["model_id"] == "policy-b" + assert isinstance(target_manager.get_adapter_state("policy-b").optimizer, SignSGD) + assert target_manager.get_adapter_session_spec("policy-b")["optimizer_config"]["type"] == "signsgd" + + +def test_adapter_coordinator_loads_checkpoint_without_placeholder_spec_mismatch(tmp_path): + source_manager = _build_manager(tmp_path / "source", optimizer_type="signsgd") + source_manager.register_adapter("policy-b", lr=0.1, initialize_fresh=True) + checkpoint_path = source_manager.save_adapter_state("policy-b")["path"] + + target_manager = _build_manager(tmp_path / "target", optimizer_type="adamw") + coordinator = AdapterCoordinator( + trainer=_CoordinatorTrainer(target_manager), + rank=0, + world_size=1, + cpu_group=None, + ) + + result = asyncio.run( + coordinator.handle_load_adapter_state( + { + "payload": AdapterStateData( + model_id="policy-b", + path=checkpoint_path, + load_optimizer=True, + ) + } + ) + ) + + assert result["model_id"] == "policy-b" + assert isinstance(target_manager.get_adapter_state("policy-b").optimizer, SignSGD) + assert target_manager.get_adapter_session_spec("policy-b")["optimizer_config"]["learning_rate"] == pytest.approx( + 0.1 + ) + + +def test_adapter_coordinator_auto_load_evicted_uses_checkpoint_session_spec(tmp_path): + target_manager = _build_manager(tmp_path / "target", optimizer_type="adamw") + source_manager = _build_manager(tmp_path / "source", optimizer_type="signsgd") + source_manager.register_adapter("policy-evicted", lr=0.2, initialize_fresh=True) + checkpoint_path = Path(target_manager.checkpoint_dir) / "evicted" / "policy-evicted" + source_manager.save_adapter_state("policy-evicted", str(checkpoint_path)) + + coordinator = AdapterCoordinator( + trainer=_CoordinatorTrainer(target_manager), + rank=0, + world_size=1, + cpu_group=None, + ) + + was_loaded, loaded_path = coordinator.auto_load_if_evicted("policy-evicted") + + assert was_loaded is True + assert loaded_path == str(checkpoint_path) + assert isinstance(target_manager.get_adapter_state("policy-evicted").optimizer, SignSGD) + assert target_manager.get_adapter_session_spec("policy-evicted")["optimizer_config"]["learning_rate"] == ( + pytest.approx(0.2) + ) + + +def test_load_adapter_state_rejects_registered_session_spec_mismatch(tmp_path): + source_manager = _build_manager(tmp_path, optimizer_type="signsgd") + source_manager.register_adapter("policy-b", lr=0.1, initialize_fresh=True) + checkpoint_path = source_manager.save_adapter_state("policy-b")["path"] + + target_manager = _build_manager(tmp_path, optimizer_type="adamw") + target_manager.register_adapter("policy-b", lr=0.1, initialize_fresh=True) + + with pytest.raises(ValueError, match="Checkpoint session spec does not match"): + target_manager.load_adapter_state("policy-b", checkpoint_path, load_optimizer=True) + + +def test_load_adapter_state_allows_lr_override_for_registered_session(tmp_path): + source_manager = _build_manager(tmp_path / "source", optimizer_type="adamw") + source_spec = _session_spec(rank=4, alpha=16, optimizer_type="adamw", lr=0.1, weight_decay=0.01) + source_manager.register_adapter("policy-lr-override", session_spec=source_spec, initialize_fresh=True) + checkpoint_path = source_manager.save_adapter_state("policy-lr-override")["path"] + + target_manager = _build_manager(tmp_path / "target", optimizer_type="adamw") + target_spec = _session_spec(rank=4, alpha=16, optimizer_type="adamw", lr=0.05, weight_decay=0.01) + target_manager.register_adapter("policy-lr-override", session_spec=target_spec, initialize_fresh=True) + + result = target_manager.load_adapter_state( + "policy-lr-override", + checkpoint_path, + load_optimizer=True, + lr=0.2, + ) + + target_state = target_manager.get_adapter_state("policy-lr-override") + assert result["model_id"] == "policy-lr-override" + assert target_state.lr == pytest.approx(0.2) + assert target_state.optimizer.param_groups[0]["lr"] == pytest.approx(0.2) + assert target_manager.get_adapter_session_spec("policy-lr-override")["optimizer_config"][ + "learning_rate" + ] == pytest.approx(0.2) + + +def test_load_adapter_state_allows_weights_only_optimizer_mismatch(tmp_path): + source_manager = _build_manager(tmp_path, optimizer_type="signsgd") + source_spec = _session_spec(rank=4, alpha=16, optimizer_type="signsgd", lr=0.1) + source_manager.register_adapter("policy-b", session_spec=source_spec, initialize_fresh=True) + source_state = source_manager.get_adapter_state("policy-b") + source_state.lora_params["model.layers.0.self_attn.o_proj.lora_A"].data.fill_(1.25) + source_state.lora_params["model.layers.0.self_attn.o_proj.lora_B"].data.fill_(0.5) + checkpoint_path = source_manager.save_adapter_state("policy-b")["path"] + + target_manager = _build_manager(tmp_path, optimizer_type="adamw") + target_spec = _session_spec(rank=4, alpha=16, optimizer_type="adamw", lr=0.05, weight_decay=0.01) + target_manager.register_adapter("policy-b", session_spec=target_spec, initialize_fresh=True) + + result = target_manager.load_adapter_state("policy-b", checkpoint_path, load_optimizer=False) + + target_state = target_manager.get_adapter_state("policy-b") + assert result["model_id"] == "policy-b" + assert isinstance(target_state.optimizer, torch.optim.AdamW) + assert target_state.lr == pytest.approx(0.05) + assert target_manager.get_adapter_session_spec("policy-b")["optimizer_config"]["type"] == "adamw" + assert target_manager.get_adapter_session_spec("policy-b")["optimizer_config"]["learning_rate"] == pytest.approx( + 0.05 + ) + assert target_state.optimizer.param_groups[0]["lr"] == pytest.approx(0.05) + assert torch.allclose( + target_state.lora_params["model.layers.0.self_attn.o_proj.lora_A"], + torch.full((4, 8), 1.25), + ) + assert torch.allclose( + target_state.lora_params["model.layers.0.self_attn.o_proj.lora_B"], + torch.full((8, 4), 0.5), + ) + + +def test_load_adapter_state_weights_only_restores_checkpoint_lr_for_same_optimizer_contract(tmp_path): + source_manager = _build_manager(tmp_path / "source", optimizer_type="adamw") + source_spec = _session_spec(rank=4, alpha=16, optimizer_type="adamw", lr=0.1, weight_decay=0.01) + source_manager.register_adapter("policy-lr-restore", session_spec=source_spec, initialize_fresh=True) + source_state = source_manager.get_adapter_state("policy-lr-restore") + for param in source_state.lora_params.values(): + param.grad = torch.ones_like(param) + source_manager.optim_step("policy-lr-restore", lr=0.25) + checkpoint_path = source_manager.save_adapter_state("policy-lr-restore")["path"] + + target_manager = _build_manager(tmp_path / "target", optimizer_type="adamw") + target_spec = _session_spec(rank=4, alpha=16, optimizer_type="adamw", lr=0.05, weight_decay=0.01) + target_manager.register_adapter("policy-lr-restore", session_spec=target_spec, initialize_fresh=True) + + target_manager.load_adapter_state("policy-lr-restore", checkpoint_path, load_optimizer=False) + + target_state = target_manager.get_adapter_state("policy-lr-restore") + assert target_state.lr == pytest.approx(0.25) + assert target_state.optimizer.param_groups[0]["lr"] == pytest.approx(0.25) + assert target_manager.get_adapter_session_spec("policy-lr-restore")["optimizer_config"][ + "learning_rate" + ] == pytest.approx(0.25) + + +def test_load_adapter_state_rejects_checkpoint_target_module_mismatch(tmp_path): + source_manager = _build_manager(tmp_path, optimizer_type="adamw") + source_manager.register_adapter("policy-structure", lr=0.1, initialize_fresh=True) + checkpoint_path = Path(source_manager.save_adapter_state("policy-structure")["path"]) + + adapter_config = json.loads((checkpoint_path / "adapter_config.json").read_text(encoding="utf-8")) + adapter_config["target_modules"] = ["q_proj"] + (checkpoint_path / "adapter_config.json").write_text(json.dumps(adapter_config), encoding="utf-8") + + target_manager = _build_manager(tmp_path, optimizer_type="adamw") + with pytest.raises(ValueError, match="target_modules"): + target_manager.load_adapter_state("policy-structure", str(checkpoint_path), load_optimizer=True) + + +def test_load_adapter_state_rejects_checkpoint_with_missing_lora_tensors(tmp_path): + source_manager = _build_manager(tmp_path, optimizer_type="adamw") + source_manager.register_adapter("policy-missing", lr=0.1, initialize_fresh=True) + checkpoint_path = Path(source_manager.save_adapter_state("policy-missing")["path"]) + + weights_path = checkpoint_path / "adapter_model.safetensors" + weights = safetensors_load_file(str(weights_path)) + weights.pop("base_model.model.model.layers.0.self_attn.o_proj.lora_B") + safetensors_save_file(weights, str(weights_path)) + + target_manager = _build_manager(tmp_path, optimizer_type="adamw") + with pytest.raises(ValueError, match="parameter set does not match"): + target_manager.load_adapter_state("policy-missing", str(checkpoint_path), load_optimizer=True) + + +def test_load_adapter_state_rolls_back_freshly_registered_adapter_on_failure(tmp_path): + source_manager = _build_manager(tmp_path, optimizer_type="adamw") + source_manager.register_adapter("policy-rollback", lr=0.1, initialize_fresh=True) + checkpoint_path = Path(source_manager.save_adapter_state("policy-rollback")["path"]) + + weights_path = checkpoint_path / "adapter_model.safetensors" + weights = safetensors_load_file(str(weights_path)) + weights.pop("base_model.model.model.layers.0.self_attn.o_proj.lora_B") + safetensors_save_file(weights, str(weights_path)) + + target_manager = _build_manager(tmp_path, optimizer_type="adamw") + assert "policy-rollback" not in target_manager.adapters + + with pytest.raises(ValueError, match="parameter set does not match"): + target_manager.load_adapter_state("policy-rollback", str(checkpoint_path), load_optimizer=True) + + assert "policy-rollback" not in target_manager.adapters + + +def test_load_adapter_state_accepts_weight_suffixed_checkpoint_tensor_names(tmp_path): + source_manager = _build_manager(tmp_path / "source", optimizer_type="adamw") + source_manager.register_adapter("policy-weight-suffix", lr=0.1, initialize_fresh=True) + source_state = source_manager.get_adapter_state("policy-weight-suffix") + source_state.lora_params["model.layers.0.self_attn.o_proj.lora_A"].data.fill_(1.5) + source_state.lora_params["model.layers.0.self_attn.o_proj.lora_B"].data.fill_(0.75) + checkpoint_path = Path(source_manager.save_adapter_state("policy-weight-suffix")["path"]) + + weights_path = checkpoint_path / "adapter_model.safetensors" + weights = safetensors_load_file(str(weights_path)) + renamed_weights = {} + for key, value in weights.items(): + renamed_weights[f"{key}.weight"] = value + safetensors_save_file(renamed_weights, str(weights_path)) + + target_manager = _build_manager(tmp_path / "target", optimizer_type="adamw") + result = target_manager.load_adapter_state("policy-weight-suffix", str(checkpoint_path), load_optimizer=True) + + target_state = target_manager.get_adapter_state("policy-weight-suffix") + assert result["model_id"] == "policy-weight-suffix" + assert torch.allclose( + target_state.lora_params["model.layers.0.self_attn.o_proj.lora_A"], + torch.full((4, 8), 1.5), + ) + assert torch.allclose( + target_state.lora_params["model.layers.0.self_attn.o_proj.lora_B"], + torch.full((8, 4), 0.75), + ) + + +def test_load_adapter_state_rejects_checkpoint_rank_exceeding_model_capacity(tmp_path): + source_manager = _build_manager( + tmp_path / "source", + optimizer_type="adamw", + max_rank=8, + lora_config={"base_model": "Qwen/Qwen3-8B", "lora_rank": 8, "lora_alpha": 16}, + ) + source_manager.register_adapter( + "policy-r8", + session_spec=_session_spec(rank=8, alpha=16, optimizer_type="adamw", lr=0.1), + initialize_fresh=True, + ) + checkpoint_path = source_manager.save_adapter_state("policy-r8")["path"] + + target_manager = _build_manager( + tmp_path / "target", + optimizer_type="adamw", + max_rank=4, + lora_config={"base_model": "Qwen/Qwen3-8B", "lora_rank": 4, "lora_alpha": 16}, + ) + with pytest.raises(ValueError, match="exceeds live model LoRA capacity"): + target_manager.load_adapter_state("policy-r8", checkpoint_path, load_optimizer=True) + + +def test_register_adapter_refuses_to_evict_dirty_adapter(tmp_path): + manager = _build_manager(tmp_path, max_adapters=1, optimizer_type="adamw") + manager.register_adapter("policy-a", session_spec=_session_spec(rank=4, alpha=16, optimizer_type="adamw", lr=0.1)) + + dirty_state = manager.get_adapter_state("policy-a") + dirty_state.lora_params["model.layers.0.self_attn.o_proj.lora_A"].grad = torch.ones_like( + dirty_state.lora_params["model.layers.0.self_attn.o_proj.lora_A"] + ) + + with pytest.raises(RuntimeError, match="pending gradients"): + manager.register_adapter( + "policy-b", session_spec=_session_spec(rank=4, alpha=16, optimizer_type="adamw", lr=0.2) + ) + + assert manager.has_adapter("policy-a") + assert not manager.has_adapter("policy-b") + + +def test_register_adapter_evicts_clean_adapter_before_dirty_one(tmp_path): + manager = _build_manager(tmp_path, max_adapters=2, optimizer_type="adamw") + manager.register_adapter("policy-a", session_spec=_session_spec(rank=4, alpha=16, optimizer_type="adamw", lr=0.1)) + manager.register_adapter("policy-b", session_spec=_session_spec(rank=4, alpha=16, optimizer_type="adamw", lr=0.2)) + + dirty_state = manager.get_adapter_state("policy-a") + clean_state = manager.get_adapter_state("policy-b") + dirty_state.last_access_time = 1.0 + clean_state.last_access_time = 2.0 + dirty_state.lora_params["model.layers.0.self_attn.o_proj.lora_A"].grad = torch.ones_like( + dirty_state.lora_params["model.layers.0.self_attn.o_proj.lora_A"] + ) + + manager.register_adapter("policy-c", session_spec=_session_spec(rank=4, alpha=16, optimizer_type="adamw", lr=0.3)) + + assert manager.has_adapter("policy-a") + assert not manager.has_adapter("policy-b") + assert manager.has_adapter("policy-c") + + +def test_multi_adapter_manager_supports_mixed_ranks_and_optimizers(tmp_path): + manager = _build_manager( + tmp_path, + optimizer_type="adamw", + lora_config={"base_model": "Qwen/Qwen3-8B", "lora_rank": 4, "lora_alpha": 16}, + ) + + small_spec = _session_spec(rank=2, alpha=8, optimizer_type="signsgd", lr=0.2) + large_spec = _session_spec(rank=4, alpha=16, optimizer_type="adamw", lr=0.05, weight_decay=0.01) + + manager.register_adapter("policy-small", session_spec=small_spec, initialize_fresh=True) + manager.register_adapter("policy-large", session_spec=large_spec, initialize_fresh=True) + + small_state = manager.get_adapter_state("policy-small") + large_state = manager.get_adapter_state("policy-large") + layer = manager.model.model.layers[0].self_attn.o_proj + + assert isinstance(small_state.optimizer, SignSGD) + assert isinstance(large_state.optimizer, torch.optim.AdamW) + assert tuple(small_state.lora_params["model.layers.0.self_attn.o_proj.lora_A"].shape) == (2, 8) + assert tuple(small_state.lora_params["model.layers.0.self_attn.o_proj.lora_B"].shape) == (8, 2) + assert tuple(large_state.lora_params["model.layers.0.self_attn.o_proj.lora_A"].shape) == (4, 8) + assert tuple(large_state.lora_params["model.layers.0.self_attn.o_proj.lora_B"].shape) == (8, 4) + + small_state.lora_params["model.layers.0.self_attn.o_proj.lora_A"].data.fill_(1.5) + small_state.lora_params["model.layers.0.self_attn.o_proj.lora_B"].data.fill_(2.5) + large_state.lora_params["model.layers.0.self_attn.o_proj.lora_A"].data.fill_(3.5) + large_state.lora_params["model.layers.0.self_attn.o_proj.lora_B"].data.fill_(4.5) + + manager.prepare_forward("policy-small") + assert layer.active_r == 2 + assert layer.active_lora_alpha == 8 + assert torch.allclose(layer.lora_A[:2], torch.full((2, 8), 1.5)) + assert torch.count_nonzero(layer.lora_A[2:]) == 0 + assert torch.allclose(layer.lora_B[:, :2], torch.full((8, 2), 2.5)) + assert torch.count_nonzero(layer.lora_B[:, 2:]) == 0 + + layer.lora_A.grad = torch.full_like(layer.lora_A, 1.0) + layer.lora_B.grad = torch.full_like(layer.lora_B, 2.0) + manager.capture_gradients("policy-small") + assert tuple(small_state.lora_params["model.layers.0.self_attn.o_proj.lora_A"].grad.shape) == (2, 8) + assert tuple(small_state.lora_params["model.layers.0.self_attn.o_proj.lora_B"].grad.shape) == (8, 2) + small_grad_norm = manager.optim_step("policy-small", lr=0.2) + assert small_grad_norm > 0 + assert manager.get_global_step("policy-small") == 1 + + manager.prepare_forward("policy-large") + assert layer.active_r == 4 + assert layer.active_lora_alpha == 16 + assert torch.allclose(layer.lora_A, torch.full((4, 8), 3.5)) + assert torch.allclose(layer.lora_B, torch.full((8, 4), 4.5)) + + layer.lora_A.grad = torch.full_like(layer.lora_A, 3.0) + layer.lora_B.grad = torch.full_like(layer.lora_B, 4.0) + manager.capture_gradients("policy-large") + assert tuple(large_state.lora_params["model.layers.0.self_attn.o_proj.lora_A"].grad.shape) == (4, 8) + assert tuple(large_state.lora_params["model.layers.0.self_attn.o_proj.lora_B"].grad.shape) == (8, 4) + large_grad_norm = manager.optim_step("policy-large", lr=0.05) + assert large_grad_norm > 0 + assert manager.get_global_step("policy-large") == 1 + assert large_state.optimizer.state + + small_checkpoint = manager.save_adapter_state("policy-small")["path"] + large_checkpoint = manager.save_adapter_state("policy-large")["path"] + + reloaded_manager = _build_manager( + tmp_path, + optimizer_type="sgd", + lora_config={"base_model": "Qwen/Qwen3-8B", "lora_rank": 4, "lora_alpha": 16}, + ) + reloaded_manager.load_adapter_state("policy-small", small_checkpoint, load_optimizer=True) + reloaded_manager.load_adapter_state("policy-large", large_checkpoint, load_optimizer=True) + + assert isinstance(reloaded_manager.get_adapter_state("policy-small").optimizer, SignSGD) + assert isinstance(reloaded_manager.get_adapter_state("policy-large").optimizer, torch.optim.AdamW) + assert reloaded_manager.get_adapter_session_spec("policy-small")["lora_config"]["lora_rank"] == 2 + assert reloaded_manager.get_adapter_session_spec("policy-large")["lora_config"]["lora_rank"] == 4 + + +def test_save_adapter_state_persists_current_learning_rate(tmp_path): + manager = _build_manager(tmp_path, optimizer_type="adamw") + session_spec = _session_spec(rank=4, alpha=16, optimizer_type="adamw", lr=0.1, weight_decay=0.01) + manager.register_adapter("policy-lr", session_spec=session_spec, initialize_fresh=True) + + state = manager.get_adapter_state("policy-lr") + for param in state.lora_params.values(): + param.grad = torch.ones_like(param) + + manager.optim_step("policy-lr", lr=0.25) + checkpoint_path = Path(manager.save_adapter_state("policy-lr")["path"]) + session_spec_json = json.loads((checkpoint_path / "session_spec.json").read_text(encoding="utf-8")) + metadata_json = json.loads((checkpoint_path / "metadata.json").read_text(encoding="utf-8")) + + assert manager.get_adapter_session_spec("policy-lr")["optimizer_config"]["learning_rate"] == pytest.approx(0.25) + assert session_spec_json["optimizer_config"]["learning_rate"] == pytest.approx(0.25) + assert metadata_json["lr"] == pytest.approx(0.25) + assert metadata_json["optimizer"]["learning_rate"] == pytest.approx(0.25) + + +def test_register_adapter_hoists_common_adam_hparams_out_of_optimizer_kwargs(tmp_path): + manager = _build_manager( + tmp_path, + optimizer_type="adamw", + lora_config={"base_model": "Qwen/Qwen3-8B", "lora_rank": 4, "lora_alpha": 16}, + ) + + session_spec = normalize_session_spec( + base_model="Qwen/Qwen3-8B", + raw_lora_config={"lora_rank": 4, "lora_alpha": 16}, + raw_optimizer_config={ + "type": "adamw", + "learning_rate": 1e-4, + "weight_decay": 0.02, + "optimizer_kwargs": { + "betas": [0.8, 0.88], + "eps": 1e-7, + "capturable": True, + }, + }, + default_rank=4, + default_alpha=16, + max_lora_rank=16, + default_optimizer_type="adamw", + default_learning_rate=1e-5, + default_weight_decay=0.01, + default_optimizer_dtype="bf16", + default_optimizer_kwargs={}, + server_lora_config={"enable_lora": True, "lora_rank": 4, "lora_alpha": 16, "max_lora_rank": 16}, + ) + + manager.register_adapter("policy-adam", session_spec=session_spec, initialize_fresh=True) + state = manager.get_adapter_state("policy-adam") + + assert isinstance(state.optimizer, torch.optim.AdamW) + assert state.optimizer.defaults["betas"] == (0.8, 0.88) + assert state.optimizer.defaults["eps"] == pytest.approx(1e-7) + assert state.optimizer.defaults["capturable"] is True + assert manager.get_adapter_session_spec("policy-adam")["optimizer_config"]["optimizer_kwargs"] == { + "capturable": True + } + + +def test_muon_set_lr_preserves_muon_param_group_lr(tmp_path): + manager = _build_manager( + tmp_path, + optimizer_type="muon", + lora_config={"base_model": "Qwen/Qwen3-8B", "lora_rank": 4, "lora_alpha": 16}, + ) + session_spec = _session_spec(rank=4, alpha=16, optimizer_type="muon", lr=1e-4, weight_decay=0.01) + session_spec["optimizer_config"]["optimizer_kwargs"] = { + "muon_lr": 0.02, + "muon_ns_use_quack_kernels": False, + } + + manager.register_adapter("policy-muon", session_spec=session_spec, initialize_fresh=True) + state = manager.get_adapter_state("policy-muon") + muon_groups = [param_group for param_group in state.optimizer.param_groups if param_group.get("use_muon", False)] + assert muon_groups + assert all(param_group["lr"] == pytest.approx(0.02) for param_group in muon_groups) + + manager.set_lr("policy-muon", 2e-4) + assert state.lr == pytest.approx(2e-4) + assert all(param_group["lr"] == pytest.approx(0.02) for param_group in muon_groups) + + manager.optim_step("policy-muon", lr=3e-4) + assert state.lr == pytest.approx(3e-4) + assert all(param_group["lr"] == pytest.approx(0.02) for param_group in muon_groups) diff --git a/tests/server/runner/test_batch_utils.py b/tests/server/runner/test_batch_utils.py new file mode 100644 index 00000000..b04499de --- /dev/null +++ b/tests/server/runner/test_batch_utils.py @@ -0,0 +1,68 @@ +from unittest.mock import Mock, patch + +import pytest +import torch + +from xorl.server.runner.utils.batch_utils import convert_batch_to_tensors, simple_sequence_shard + + +pytestmark = [pytest.mark.cpu, pytest.mark.server] + + +def test_convert_batch_to_tensors_preserves_teacher_hidden_state_floats(): + batch = { + "teacher_hidden_states": [ + [[0.25, -1.75], [2.5, 3.125]], + ], + } + + converted = convert_batch_to_tensors(batch) + + assert converted["teacher_hidden_states"].dtype == torch.float32 + assert converted["teacher_hidden_states"].shape == (1, 2, 2) + torch.testing.assert_close( + converted["teacher_hidden_states"], + torch.tensor([[[0.25, -1.75], [2.5, 3.125]]], dtype=torch.float32), + ) + + +def test_convert_batch_to_tensors_pads_ragged_teacher_hidden_states(): + batch = { + "teacher_hidden_states": [ + [[0.25, 0.5]], + [[1.25, 1.5], [2.25, 2.5]], + ], + } + + converted = convert_batch_to_tensors(batch) + + assert converted["teacher_hidden_states"].shape == (2, 2, 2) + torch.testing.assert_close( + converted["teacher_hidden_states"], + torch.tensor( + [ + [[0.25, 0.5], [0.0, 0.0]], + [[1.25, 1.5], [2.25, 2.5]], + ], + dtype=torch.float32, + ), + ) + + +@patch("xorl.server.runner.utils.batch_utils.get_parallel_state") +def test_simple_sequence_shard_slices_teacher_hidden_states_on_sequence_dim(mock_parallel_state): + mock_parallel_state.return_value = Mock(cp_size=2, cp_rank=1) + batch = { + "input_ids": torch.tensor([[1, 2, 3]]), + "labels": torch.tensor([[2, 3, -100]]), + "position_ids": torch.tensor([[0, 1, 2]]), + "teacher_hidden_states": torch.tensor([[[0.25, 0.5], [1.25, 1.5], [2.25, 2.5]]]), + } + + sharded = simple_sequence_shard(batch) + + assert sharded["teacher_hidden_states"].shape == (1, 2, 2) + torch.testing.assert_close( + sharded["teacher_hidden_states"], + torch.tensor([[[2.25, 2.5], [0.0, 0.0]]]), + ) diff --git a/tests/server/runner/test_checkpoint_loading.py b/tests/server/runner/test_checkpoint_loading.py new file mode 100644 index 00000000..cde38972 --- /dev/null +++ b/tests/server/runner/test_checkpoint_loading.py @@ -0,0 +1,87 @@ +import pytest +import torch +from torch import nn + +from xorl.server.runner.checkpoint.manager import CheckpointManager +from xorl.server.runner.model_runner import ModelRunner + + +pytestmark = [pytest.mark.cpu, pytest.mark.server] + + +class _MetaModel(nn.Module): + def __init__(self): + super().__init__() + self.weight = nn.Parameter(torch.empty(2, device="meta")) + self.to_empty_device = None + + def to_empty(self, *, device, recurse=True): # noqa: ARG002 + self.to_empty_device = device + self.weight = nn.Parameter(torch.empty(2, device="cpu")) + return self + + +class _DummyCheckpointer: + def __init__(self): + self.path = None + self.state = None + + def load(self, path, state): + self.path = path + self.state = state + state["extra_state"].update( + { + "global_step": 7, + "global_forward_backward_step": 11, + "torch_rng_state": torch.get_rng_state(), + } + ) + + +def test_checkpoint_manager_materializes_skip_mode_and_omits_missing_optimizer(monkeypatch): + model = _MetaModel() + checkpointer = _DummyCheckpointer() + manager = CheckpointManager( + model=model, + optimizer=object(), + checkpointer=checkpointer, + lora_config={}, + model_config={}, + train_config={"load_weights_mode": "skip"}, + rank=0, + local_rank=0, + ) + monkeypatch.setattr("xorl.server.runner.checkpoint.manager.get_device_type", lambda: "cpu") + monkeypatch.setattr(manager, "_checkpoint_has_optimizer", lambda _path: False) + + result = manager.load_state("/tmp/model-only-dcp", load_optimizer=True) + + assert model.to_empty_device == "cpu" + assert checkpointer.path == "/tmp/model-only-dcp" + assert "optimizer" not in checkpointer.state + assert result["load_optimizer"] is False + assert manager.global_step == 7 + assert manager.global_forward_backward_step == 11 + + +def test_model_runner_loads_initial_checkpoint_and_syncs_state(): + class FakeCheckpointManager: + def __init__(self): + self.calls = [] + self.global_step = 13 + self.global_forward_backward_step = 17 + + def load_state(self, checkpoint_path, load_optimizer=True): + self.calls.append((checkpoint_path, load_optimizer)) + + runner = object.__new__(ModelRunner) + runner.train_config = {"load_checkpoint_path": "/tmp/initial-dcp"} + runner.global_step = 0 + runner.global_forward_backward_step = 0 + runner._checkpoint_mgr = FakeCheckpointManager() + + runner._load_initial_checkpoint() + + assert runner._checkpoint_mgr.calls == [("/tmp/initial-dcp", True)] + assert runner.global_step == 13 + assert runner.global_forward_backward_step == 17 diff --git a/tests/server/runner/test_checkpoint_manager_save_failures.py b/tests/server/runner/test_checkpoint_manager_save_failures.py new file mode 100644 index 00000000..e81d66eb --- /dev/null +++ b/tests/server/runner/test_checkpoint_manager_save_failures.py @@ -0,0 +1,336 @@ +"""Tests for checkpoint-manager save failure handling.""" + +import importlib.util +import json +from pathlib import Path + +import pytest +import torch +import torch.nn as nn + +from xorl.server.runner.adapters.manager import LoRAAdapterManager +from xorl.server.session_spec import normalize_session_spec + + +_MODULE_PATH = Path(__file__).resolve().parents[3] / "src" / "xorl" / "server" / "runner" / "checkpoint" / "manager.py" +_SPEC = importlib.util.spec_from_file_location("xorl_test_checkpoint_manager", _MODULE_PATH) +assert _SPEC is not None and _SPEC.loader is not None +_MODULE = importlib.util.module_from_spec(_SPEC) +_SPEC.loader.exec_module(_MODULE) +CheckpointManager = _MODULE.CheckpointManager + + +pytestmark = [pytest.mark.cpu, pytest.mark.server] + + +class _FakeOptimizer: + def state_dict(self): + return {"state": {}, "param_groups": []} + + +class _FakeAdapterState: + def __init__(self): + self.global_step = 7 + self.global_forward_backward_step = 11 + self.lr = 2e-5 + self.optimizer = _FakeOptimizer() + self.lora_params = {"adapter.weight.lora_A": nn.Parameter(torch.ones(1, 1))} + self.session_spec = { + "lora_config": { + "lora_rank": 4, + "lora_alpha": 16, + }, + "optimizer_config": { + "type": "adamw", + "learning_rate": 2e-5, + "weight_decay": 0.01, + "optimizer_dtype": "bf16", + "betas": [0.9, 0.95], + "eps": 1e-8, + "optimizer_kwargs": {}, + }, + } + + +class _FakeAdapterManager: + def __init__(self): + self.checkpoint_dir = "/tmp/adapters" + self.current_adapter_id = "policy-a" + self.adapters = {"policy-a": _FakeAdapterState()} + + def get_adapter_state(self, model_id: str): + return self.adapters[model_id] + + def get_global_step(self, model_id: str) -> int: + return self.adapters[model_id].global_step + + def get_adapter_session_spec(self, model_id: str): + return self.adapters[model_id].session_spec + + def switch_adapter(self, model_id: str, auto_register: bool = False) -> bool: + return model_id in self.adapters + + +class _DummyLoRALayer(nn.Module): + def __init__(self, *, max_rank: int = 4) -> None: + super().__init__() + self.lora_A = nn.Parameter(torch.randn(max_rank, 8)) + self.lora_B = nn.Parameter(torch.zeros(8, max_rank)) + self.active_r = max_rank + self.active_lora_alpha = 16 + + def set_runtime_lora_config(self, lora_rank: int, lora_alpha: int) -> None: + self.active_r = lora_rank + self.active_lora_alpha = lora_alpha + + +class _DummyLoRAModel(nn.Module): + def __init__(self, *, max_rank: int = 4) -> None: + super().__init__() + self.model = nn.Module() + self.model.layers = nn.ModuleList([nn.Module()]) + self.model.layers[0].self_attn = nn.Module() + self.model.layers[0].self_attn.o_proj = _DummyLoRALayer(max_rank=max_rank) + + +def _build_checkpoint_manager() -> CheckpointManager: + manager = object.__new__(CheckpointManager) + manager.rank = 0 + manager.local_rank = 0 + manager.lora_config = {"enable_lora": True} + manager._adapter_manager = _FakeAdapterManager() + return manager + + +def _build_fast_save_manager(tmp_path: Path) -> CheckpointManager: + model = _DummyLoRAModel(max_rank=4) + adapter_manager = LoRAAdapterManager( + model, + device=torch.device("cpu"), + checkpoint_dir=str(tmp_path / "adapters"), + auto_save_on_eviction=False, + lora_config={ + "base_model": "Qwen/Qwen3-8B", + "lora_rank": 4, + "lora_alpha": 16, + }, + ) + adapter_manager.register_adapter( + "policy-a", + session_spec=normalize_session_spec( + base_model="Qwen/Qwen3-8B", + raw_lora_config={ + "lora_rank": 4, + "lora_alpha": 16, + }, + raw_optimizer_config={ + "type": "adamw", + "learning_rate": 1e-4, + "weight_decay": 0.01, + "optimizer_dtype": "bf16", + "betas": [0.9, 0.95], + "eps": 1e-8, + "optimizer_kwargs": {}, + }, + default_rank=4, + default_alpha=16, + max_lora_rank=4, + default_optimizer_type="adamw", + default_learning_rate=1e-4, + default_weight_decay=0.01, + default_optimizer_dtype="bf16", + default_optimizer_kwargs={}, + ), + initialize_fresh=True, + ) + + manager = object.__new__(CheckpointManager) + manager.rank = 0 + manager.local_rank = 0 + manager.model = model + manager.model_config = {"model_path": "Qwen/Qwen3-8B"} + manager.lora_config = { + "enable_lora": True, + "lora_rank": 4, + "lora_alpha": 16, + "lora_target_modules": ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"], + } + manager._adapter_manager = adapter_manager + return manager + + +def test_save_adapter_state_raises_before_barrier_when_rank0_write_fails(monkeypatch, tmp_path): + manager = _build_checkpoint_manager() + manager._save_lora_weights = lambda path, model_id, **kwargs: None + manager._sync_collective_error = lambda error: error + + def _fail_write(*args, **kwargs): + raise PermissionError("disk full") + + monkeypatch.setattr(manager, "_write_adapter_training_artifacts", _fail_write) + monkeypatch.setattr( + _MODULE.dist, "barrier", lambda: (_ for _ in ()).throw(AssertionError("barrier should not run")) + ) + + with pytest.raises(RuntimeError, match="Adapter state save failed: disk full"): + manager.save_adapter_state("policy-a", path=str(tmp_path / "adapter-save"), save_optimizer=True) + + +def test_save_adapter_state_requests_dtype_preserving_lora_checkpoint(monkeypatch, tmp_path): + manager = _build_checkpoint_manager() + captured = {} + + def _capture_save_lora_weights(path, model_id, **kwargs): + captured["path"] = path + captured["model_id"] = model_id + captured.update(kwargs) + + monkeypatch.setattr(manager, "_save_lora_weights", _capture_save_lora_weights) + monkeypatch.setattr(manager, "_write_adapter_training_artifacts", lambda *args, **kwargs: None) + + manager.save_adapter_state("policy-a", path=str(tmp_path / "adapter-save"), save_optimizer=True) + + assert captured["model_id"] == "policy-a" + assert captured["preserve_lora_dtype"] is True + + +def test_save_lora_only_raises_before_barrier_when_rank0_write_fails(monkeypatch, tmp_path): + manager = _build_checkpoint_manager() + manager._sync_collective_error = lambda error: error + + def _fail_save(*args, **kwargs): + raise PermissionError("peft write failed") + + monkeypatch.setattr(manager, "_save_lora_weights", _fail_save) + monkeypatch.setattr( + _MODULE.dist, "barrier", lambda: (_ for _ in ()).throw(AssertionError("barrier should not run")) + ) + + with pytest.raises(RuntimeError, match="LoRA-only save failed: peft write failed"): + manager.save_lora_only(str(tmp_path / "adapter-export"), model_id="policy-a") + + +def test_fast_lora_save_uses_live_adapter_target_modules_not_requested_config(tmp_path): + manager = _build_fast_save_manager(tmp_path) + + export_dir = tmp_path / "adapter-export" + manager._save_lora_weights(str(export_dir), "policy-a") + + adapter_config = json.loads((export_dir / "adapter_config.json").read_text(encoding="utf-8")) + + assert sorted(adapter_config["target_modules"]) == ["o_proj"] + + +def test_moe_lora_save_uses_collective_gather_even_with_adapter_manager(monkeypatch, tmp_path): + manager = _build_checkpoint_manager() + manager.model = nn.Module() + manager.model_config = {"model_path": "Qwen/Qwen3-8B"} + manager.lora_config = { + "enable_lora": True, + "moe_hybrid_shared_lora": True, + "lora_rank": 4, + "lora_alpha": 16, + "lora_target_modules": ["gate_proj", "up_proj", "down_proj"], + } + manager._gather_adapter_lora_params = lambda model_id: (_ for _ in ()).throw( + AssertionError("fast adapter-manager gather should not run for MoE LoRA") + ) + + collective_state = { + "model.layers.0.mlp.experts.gate_proj_lora_A": torch.arange(64, dtype=torch.float32).reshape(1, 8, 8), + "model.layers.0.mlp.experts.gate_proj_lora_B": torch.arange(128, dtype=torch.float32).reshape(1, 8, 16), + } + monkeypatch.setattr(_MODULE, "get_lora_state_dict", lambda model: collective_state) + + captured = {} + + def _capture_save_lora_checkpoint(**kwargs): + captured.update(kwargs) + + monkeypatch.setattr(_MODULE, "save_lora_checkpoint", _capture_save_lora_checkpoint) + + manager._save_lora_weights(str(tmp_path / "moe-export"), "policy-a") + + exported_state = captured["lora_state_dict"] + assert tuple(exported_state["model.layers.0.mlp.experts.gate_proj_lora_A"].shape) == (1, 8, 4) + assert tuple(exported_state["model.layers.0.mlp.experts.gate_proj_lora_B"].shape) == (1, 4, 16) + torch.testing.assert_close( + exported_state["model.layers.0.mlp.experts.gate_proj_lora_A"], + collective_state["model.layers.0.mlp.experts.gate_proj_lora_A"][..., :4], + ) + torch.testing.assert_close( + exported_state["model.layers.0.mlp.experts.gate_proj_lora_B"], + collective_state["model.layers.0.mlp.experts.gate_proj_lora_B"][:, :4, :], + ) + assert captured["r"] == 4 + + +def test_moe_lora_save_uses_resolved_target_modules_for_detection(monkeypatch, tmp_path): + manager = _build_checkpoint_manager() + + class _ModelWithStackedMoELoRA: + def named_parameters(self): + yield "model.layers.0.mlp.experts.gate_proj_lora_A", nn.Parameter(torch.ones(1, 8, 4)) + + manager.model = _ModelWithStackedMoELoRA() + manager.model_config = {"model_path": "Qwen/Qwen3-8B"} + manager.lora_config = { + "enable_lora": True, + "lora_rank": 4, + "lora_alpha": 16, + } + manager.lora_target_modules = ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"] + manager.lora_alpha_value = 16 + manager._gather_adapter_lora_params = lambda model_id: (_ for _ in ()).throw( + AssertionError("fast adapter-manager gather should not run for resolved MoE LoRA targets") + ) + + collective_state = {"model.layers.0.mlp.experts.gate_proj_lora_A": torch.ones(1, 8, 4)} + monkeypatch.setattr(_MODULE, "get_lora_state_dict", lambda model: collective_state) + + captured = {} + + def _capture_save_lora_checkpoint(**kwargs): + captured.update(kwargs) + + monkeypatch.setattr(_MODULE, "save_lora_checkpoint", _capture_save_lora_checkpoint) + + manager._save_lora_weights(str(tmp_path / "moe-export"), "policy-a") + + assert captured["lora_state_dict"].keys() == collective_state.keys() + torch.testing.assert_close( + captured["lora_state_dict"]["model.layers.0.mlp.experts.gate_proj_lora_A"], + collective_state["model.layers.0.mlp.experts.gate_proj_lora_A"], + ) + assert captured["r"] == 4 + + +def test_lora_save_forwards_export_format(monkeypatch, tmp_path): + manager = object.__new__(CheckpointManager) + manager.rank = 0 + manager.local_rank = 0 + manager.model = nn.Module() + manager.model_config = {"model_path": "Qwen/Qwen3-8B"} + manager._adapter_manager = None + manager.lora_target_modules = ["gate_proj", "up_proj", "down_proj"] + manager.lora_alpha_value = 16 + manager.lora_config = { + "enable_lora": True, + "moe_hybrid_shared_lora": True, + "lora_rank": 4, + "lora_alpha": 16, + "lora_export_format": "sglang_shared_outer", + } + + monkeypatch.setattr(_MODULE, "get_lora_state_dict", lambda model: {}) + + captured = {} + + def _capture_save_lora_checkpoint(**kwargs): + captured.update(kwargs) + + monkeypatch.setattr(_MODULE, "save_lora_checkpoint", _capture_save_lora_checkpoint) + + manager._save_lora_weights(str(tmp_path / "sglang-export"), "default") + + assert captured["lora_export_format"] == "sglang_shared_outer" diff --git a/tests/server/runner/test_deepseek_v3_lora_targets.py b/tests/server/runner/test_deepseek_v3_lora_targets.py new file mode 100644 index 00000000..559c5d4c --- /dev/null +++ b/tests/server/runner/test_deepseek_v3_lora_targets.py @@ -0,0 +1,102 @@ +import json + +import pytest + +from xorl.server.runner.model_runner import ModelRunner + + +pytestmark = [pytest.mark.cpu, pytest.mark.server] + + +def _write_kimi_config(config_dir): + config_dir.mkdir() + (config_dir / "config.json").write_text( + json.dumps( + { + "model_type": "kimi_k25", + "text_config": { + "model_type": "kimi_k2", + "architectures": ["DeepseekV3ForCausalLM"], + "vocab_size": 163840, + "hidden_size": 128, + "intermediate_size": 256, + "moe_intermediate_size": 32, + "num_hidden_layers": 4, + "num_attention_heads": 4, + "num_key_value_heads": 4, + "kv_lora_rank": 16, + "q_lora_rank": 32, + "qk_nope_head_dim": 16, + "qk_rope_head_dim": 8, + "v_head_dim": 16, + "n_routed_experts": 8, + "n_shared_experts": 2, + "n_group": 4, + "topk_group": 2, + "num_experts_per_tok": 2, + "first_k_dense_replace": 1, + "routed_scaling_factor": 1.25, + "norm_topk_prob": True, + "hidden_act": "silu", + "max_position_embeddings": 4096, + "initializer_range": 0.02, + "rms_norm_eps": 1e-6, + "use_cache": True, + "attention_bias": False, + "attention_dropout": 0.0, + "output_router_logits": True, + "router_aux_loss_coef": 0.001, + "topk_method": "noaux_tc", + "scoring_func": "sigmoid", + "rope_scaling": {"rope_type": "default", "rope_theta": 1000000.0}, + }, + "vision_config": {"hidden_size": 1024}, + "tie_word_embeddings": False, + } + ) + ) + + +def _make_runner(config_path, lora_config): + runner = object.__new__(ModelRunner) + runner.model_config = {"config_path": str(config_path)} + runner.lora_config = lora_config + return runner + + +def test_model_runner_resolves_kimi_defaults_from_wrapper_config(tmp_path): + config_dir = tmp_path / "kimi-k25" + _write_kimi_config(config_dir) + runner = _make_runner( + config_dir, + { + "train_attn": True, + "train_mlp": False, + "train_unembed": True, + }, + ) + + assert runner._resolve_lora_target_modules() == [ + "q_a_proj", + "q_b_proj", + "kv_a_proj_with_mqa", + "kv_b_proj", + "o_proj", + "lm_head", + ] + + +def test_model_runner_prefers_explicit_lora_targets(tmp_path): + config_dir = tmp_path / "kimi-k25" + _write_kimi_config(config_dir) + runner = _make_runner( + config_dir, + { + "lora_target_modules": ["q_b_proj", "o_proj"], + "train_attn": False, + "train_mlp": False, + "train_unembed": False, + }, + ) + + assert runner._resolve_lora_target_modules() == ["q_b_proj", "o_proj"] diff --git a/tests/server/runner/test_lora_checkpoint_roundtrip.py b/tests/server/runner/test_lora_checkpoint_roundtrip.py index 13c1f9fd..ec1b45fe 100644 --- a/tests/server/runner/test_lora_checkpoint_roundtrip.py +++ b/tests/server/runner/test_lora_checkpoint_roundtrip.py @@ -1,13 +1,22 @@ import importlib.util +import json from pathlib import Path import pytest import torch import torch.nn as nn +from safetensors.torch import load_file as load_safetensors_file from xorl.lora.modules.linear import LoraLinear -from xorl.lora.utils import load_lora_checkpoint, save_lora_checkpoint +from xorl.lora.utils import ( + LoraTensorShardSpec, + convert_peft_lora_state_dict, + get_lora_state_dict, + load_lora_checkpoint, + save_lora_checkpoint, +) from xorl.models.layers.moe import MoEExpertsLoRA, MoELoRAConfig +from xorl.optim import AnyPrecisionAdamW pytestmark = [pytest.mark.cpu, pytest.mark.server] @@ -25,35 +34,35 @@ class _TinyAttention(nn.Module): - def __init__(self): + def __init__(self, r: int = 2, lora_alpha: int = 4): super().__init__() - self.q_proj = LoraLinear.from_module(nn.Linear(8, 8, bias=False), r=2, lora_alpha=4) + self.q_proj = LoraLinear.from_module(nn.Linear(8, 8, bias=False), r=r, lora_alpha=lora_alpha) class _TinyLayer(nn.Module): - def __init__(self): + def __init__(self, r: int = 2, lora_alpha: int = 4): super().__init__() - self.self_attn = _TinyAttention() + self.self_attn = _TinyAttention(r=r, lora_alpha=lora_alpha) 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), + lora_config=MoELoRAConfig(r=r, lora_alpha=lora_alpha, hybrid_shared=True), ) class _TinyInnerModel(nn.Module): - def __init__(self): + def __init__(self, r: int = 2, lora_alpha: int = 4): super().__init__() - self.layers = nn.ModuleList([_TinyLayer()]) + self.layers = nn.ModuleList([_TinyLayer(r=r, lora_alpha=lora_alpha)]) class _TinyMoELoraModel(nn.Module): - def __init__(self): + def __init__(self, r: int = 2, lora_alpha: int = 4): super().__init__() - self.model = _TinyInnerModel() + self.model = _TinyInnerModel(r=r, lora_alpha=lora_alpha) def _iter_lora_parameters(module: nn.Module): @@ -80,6 +89,127 @@ 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)} +@pytest.mark.parametrize( + ("proj_name", "lora_type", "per_expert_shape"), [("down_proj", "A", (5, 2)), ("gate_proj", "B", (2, 5))] +) +def test_convert_peft_moe_lora_slices_global_experts_for_ep_shard(proj_name, lora_type, per_expert_shape): + prefix = "model.layers.0" + internal_name = f"{prefix}.mlp.experts.{proj_name}_lora_{lora_type}" + global_tensor = torch.arange(8 * per_expert_shape[0] * per_expert_shape[1], dtype=torch.float32).reshape( + 8, *per_expert_shape + ) + checkpoint_state = { + f"base_model.model.{prefix}.mlp.experts.{expert_idx}.{proj_name}.lora_{lora_type}.weight": global_tensor[ + expert_idx + ] + .transpose(0, 1) + .contiguous() + for expert_idx in range(8) + } + + converted = convert_peft_lora_state_dict( + checkpoint_state, + expected_shapes={internal_name: torch.Size((2, *per_expert_shape))}, + expected_shard_specs={internal_name: LoraTensorShardSpec(dim=0, index=2, size=4)}, + ) + + assert set(converted) == {internal_name} + assert torch.equal(converted[internal_name], global_tensor[4:6]) + + +def test_runtime_rank_lora_export_slices_weights_and_config(tmp_path): + source = _TinyMoELoraModel(r=4, lora_alpha=8) + _assign_distinct_lora_values(source) + for module in source.modules(): + setter = getattr(module, "set_runtime_lora_config", None) + if callable(setter): + setter(lora_rank=2, lora_alpha=6) + + state = get_lora_state_dict(source) + assert state["model.layers.0.self_attn.q_proj.lora_A"].shape == (2, 8) + assert state["model.layers.0.self_attn.q_proj.lora_B"].shape == (8, 2) + assert state["model.layers.0.mlp.experts.gate_proj_lora_A"].shape == (1, 8, 2) + assert state["model.layers.0.mlp.experts.gate_proj_lora_B"].shape == (4, 2, 16) + assert state["model.layers.0.mlp.experts.down_proj_lora_A"].shape == (4, 16, 2) + assert state["model.layers.0.mlp.experts.down_proj_lora_B"].shape == (1, 2, 8) + + checkpoint_dir = tmp_path / "checkpoint" + save_lora_checkpoint( + model=source, + save_path=str(checkpoint_dir), + target_modules=_TARGET_MODULES, + r=4, + lora_alpha=8, + moe_hybrid_shared_lora=True, + ) + + weights = load_safetensors_file(str(checkpoint_dir / "adapter_model.safetensors")) + cfg = json.loads((checkpoint_dir / "adapter_config.json").read_text()) + prefix = "base_model.model.model.layers.0" + + assert cfg["r"] == 2 + assert cfg["lora_alpha"] == 6 + assert weights[f"{prefix}.self_attn.q_proj.lora_A.weight"].shape == (2, 8) + assert weights[f"{prefix}.self_attn.q_proj.lora_B.weight"].shape == (8, 2) + assert weights[f"{prefix}.mlp.experts.shared.gate_proj.lora_A.weight"].shape == (2, 8) + assert weights[f"{prefix}.mlp.experts.0.gate_proj.lora_B.weight"].shape == (16, 2) + assert weights[f"{prefix}.mlp.experts.0.down_proj.lora_A.weight"].shape == (2, 16) + assert weights[f"{prefix}.mlp.experts.shared.down_proj.lora_B.weight"].shape == (8, 2) + + +def test_save_lora_checkpoint_exports_hybrid_shared_moe_in_peft_orientation(tmp_path): + source = _TinyMoELoraModel() + _assign_distinct_lora_values(source) + + checkpoint_dir = tmp_path / "checkpoint" + save_lora_checkpoint( + model=source, + save_path=str(checkpoint_dir), + moe_hybrid_shared_lora=True, + ) + + weights = load_safetensors_file(str(checkpoint_dir / "adapter_model.safetensors")) + with open(checkpoint_dir / "adapter_config.json", "r") as f: + adapter_config = json.load(f) + + moe = source.model.layers[0].mlp.experts + + gate_proj_shared_a = weights["base_model.model.model.layers.0.mlp.experts.shared.gate_proj.lora_A.weight"] + up_proj_shared_a = weights["base_model.model.model.layers.0.mlp.experts.shared.up_proj.lora_A.weight"] + gate_proj_expert_b = weights["base_model.model.model.layers.0.mlp.experts.0.gate_proj.lora_B.weight"] + down_proj_expert_a = weights["base_model.model.model.layers.0.mlp.experts.0.down_proj.lora_A.weight"] + down_proj_shared_b = weights["base_model.model.model.layers.0.mlp.experts.shared.down_proj.lora_B.weight"] + + assert gate_proj_shared_a.shape == (2, 8) + assert up_proj_shared_a.shape == (2, 8) + assert gate_proj_expert_b.shape == (16, 2) + assert down_proj_expert_a.shape == (2, 16) + assert down_proj_shared_b.shape == (8, 2) + + assert torch.equal( + gate_proj_shared_a, + moe.gate_proj_lora_A.detach().cpu()[0].transpose(0, 1).contiguous().to(torch.bfloat16), + ) + assert torch.equal( + up_proj_shared_a, + moe.up_proj_lora_A.detach().cpu()[0].transpose(0, 1).contiguous().to(torch.bfloat16), + ) + assert torch.equal( + gate_proj_expert_b, + moe.gate_proj_lora_B.detach().cpu()[0].transpose(0, 1).contiguous().to(torch.bfloat16), + ) + assert torch.equal( + down_proj_expert_a, + moe.down_proj_lora_A.detach().cpu()[0].transpose(0, 1).contiguous().to(torch.bfloat16), + ) + assert torch.equal( + down_proj_shared_b, + moe.down_proj_lora_B.detach().cpu()[0].transpose(0, 1).contiguous().to(torch.bfloat16), + ) + assert adapter_config["r"] == 2 + assert adapter_config["moe_hybrid_shared_lora"] is True + + def test_load_lora_checkpoint_roundtrip_supports_hybrid_shared(tmp_path): source = _TinyMoELoraModel() _assign_distinct_lora_values(source) @@ -105,6 +235,83 @@ def test_load_lora_checkpoint_roundtrip_supports_hybrid_shared(tmp_path): assert torch.equal(actual[name], expected_tensor), name +def test_save_lora_checkpoint_sglang_shared_outer_layout(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, + lora_export_format="sglang_shared_outer", + ) + + cfg = json.loads((checkpoint_dir / "adapter_config.json").read_text()) + assert cfg["_sglang_lora_format"] == "shared_outer" + assert cfg["moe_hybrid_shared_lora"] is True + + tensors = load_safetensors_file(str(checkpoint_dir / "adapter_model.safetensors")) + + # Expected shapes for the tiny hybrid-shared MoE: E=4, hidden=8, inter=16, r=2. + layer_prefix = "base_model.model.model.layers.0.mlp.experts" + expected_shapes = { + f"{layer_prefix}.w1.lora_A.weight": (1, 2, 8), + f"{layer_prefix}.w1.lora_B.weight": (4, 16, 2), + f"{layer_prefix}.w2.lora_A.weight": (4, 2, 16), + f"{layer_prefix}.w2.lora_B.weight": (1, 8, 2), + f"{layer_prefix}.w3.lora_A.weight": (1, 2, 8), + f"{layer_prefix}.w3.lora_B.weight": (4, 16, 2), + } + moe_keys = {k for k in tensors if ".mlp.experts." in k} + assert moe_keys == set(expected_shapes) + for key, shape in expected_shapes.items(): + assert tuple(tensors[key].shape) == shape, key + + +def test_save_and_load_sglang_shared_outer_hybrid_shared_roundtrip(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, + lora_export_format="sglang_shared_outer", + ) + + 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_save_lora_checkpoint_sglang_shared_outer_requires_hybrid_shared(tmp_path): + source = _TinyMoELoraModel() + with pytest.raises(ValueError, match="moe_hybrid_shared_lora=True"): + save_lora_checkpoint( + model=source, + save_path=str(tmp_path / "checkpoint"), + target_modules=_TARGET_MODULES, + r=2, + lora_alpha=4, + moe_hybrid_shared_lora=False, + lora_export_format="sglang_shared_outer", + ) + + def test_adapter_manager_load_adapter_state_roundtrip_supports_hybrid_shared(tmp_path): source = _TinyMoELoraModel() _assign_distinct_lora_values(source) @@ -142,3 +349,21 @@ def test_adapter_manager_load_adapter_state_roundtrip_supports_hybrid_shared(tmp assert set(actual) == set(expected) for name, expected_tensor in expected.items(): assert torch.equal(actual[name], expected_tensor), name + + +def test_adapter_manager_threads_cautious_weight_decay_to_optimizer(tmp_path): + 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}, + optimizer_config={"cautious_weight_decay": True, "weight_decay": 0.02}, + ) + + manager.register_adapter("adapter-cwd", lr=1e-4, initialize_fresh=True) + + optimizer = manager.get_adapter_state("adapter-cwd").optimizer + assert isinstance(optimizer, AnyPrecisionAdamW) + assert all(group.get("cautious") is True for group in optimizer.param_groups) + assert [group["weight_decay"] for group in optimizer.param_groups] == [0.02] diff --git a/tests/server/runner/test_model_runner_is_metrics.py b/tests/server/runner/test_model_runner_is_metrics.py new file mode 100644 index 00000000..e8f0edea --- /dev/null +++ b/tests/server/runner/test_model_runner_is_metrics.py @@ -0,0 +1,84 @@ +from types import SimpleNamespace +from unittest.mock import patch + +import pytest +import torch + +from xorl.server.runner import model_runner as mr +from xorl.server.runner.model_runner import ModelRunner + + +pytestmark = [pytest.mark.cpu, pytest.mark.server] + + +def test_is_metric_accumulation_outputs_python_scalars_without_dp(): + accumulated = {} + + ModelRunner._accumulate_is_metrics( + accumulated, + { + "valid_tokens": 2, + "ratio_mean": torch.tensor(3.0), + "ratio_min": torch.tensor(0.75), + "ratio_max": torch.tensor(1.25), + }, + ) + ModelRunner._accumulate_is_metrics( + accumulated, + { + "valid_tokens": 3, + "ratio_mean": torch.tensor(6.0), + "ratio_min": torch.tensor(0.5), + "ratio_max": torch.tensor(1.5), + }, + ) + + for metric in accumulated.values(): + assert isinstance(metric["sum"], float) + assert isinstance(metric["count"], (float, int)) + + result = {} + with patch.object(mr, "get_parallel_state", lambda: SimpleNamespace(dp_enabled=False)): + ModelRunner._finalize_is_metrics(accumulated, result) + + assert result == { + "is_valid_tokens": 2.5, + "is_ratio_mean": 1.8, + "is_ratio_min": 0.5, + "is_ratio_max": 1.5, + } + assert all(not isinstance(value, torch.Tensor) for value in result.values()) + + +def test_metric_ops_preserve_tis_extrema_without_dp(): + accumulated = {} + metric_ops = {"tis_min": "min", "tis_max": "max"} + + ModelRunner._accumulate_is_metrics( + accumulated, + { + "valid_tokens": 2, + "tis_mean": torch.tensor(2.0), + "tis_min": torch.tensor(0.75), + "tis_max": torch.tensor(1.25), + }, + metric_ops, + ) + ModelRunner._accumulate_is_metrics( + accumulated, + { + "valid_tokens": 3, + "tis_mean": torch.tensor(3.0), + "tis_min": torch.tensor(0.5), + "tis_max": torch.tensor(1.5), + }, + metric_ops, + ) + + result = {} + with patch.object(mr, "get_parallel_state", lambda: SimpleNamespace(dp_enabled=False)): + ModelRunner._finalize_is_metrics(accumulated, result) + + assert result["is_tis_mean"] == 1.0 + assert result["is_tis_min"] == 0.5 + assert result["is_tis_max"] == 1.5 diff --git a/tests/server/runner/test_model_runner_kill_session.py b/tests/server/runner/test_model_runner_kill_session.py new file mode 100644 index 00000000..2fd2b007 --- /dev/null +++ b/tests/server/runner/test_model_runner_kill_session.py @@ -0,0 +1,76 @@ +"""Tests for LoRA kill-session checkpoint safety in ModelRunner.""" + +import importlib.util +from pathlib import Path + +import pytest + + +_MODULE_PATH = Path(__file__).resolve().parents[3] / "src" / "xorl" / "server" / "runner" / "model_runner.py" +_SPEC = importlib.util.spec_from_file_location("xorl_test_model_runner_kill_session", _MODULE_PATH) +assert _SPEC is not None and _SPEC.loader is not None +_MODULE = importlib.util.module_from_spec(_SPEC) +_SPEC.loader.exec_module(_MODULE) +ModelRunner = _MODULE.ModelRunner + + +pytestmark = [pytest.mark.cpu, pytest.mark.server] + + +class _FakeAdapterManager: + def __init__(self, has_adapter: bool): + self._has_adapter = has_adapter + self.removed = [] + + def has_adapter(self, model_id: str) -> bool: + return self._has_adapter + + def remove_adapter(self, model_id: str) -> None: + self.removed.append(model_id) + + +def _build_runner(tmp_path: Path, *, has_adapter: bool): + runner = object.__new__(ModelRunner) + runner.rank = 0 + runner.lora_config = {"enable_lora": True} + runner.train_config = {"output_dir": str(tmp_path)} + runner._adapter_manager = _FakeAdapterManager(has_adapter=has_adapter) + runner._lora_session_specs = { + "policy-a": { + "base_model": "Qwen/Qwen3-8B", + "is_lora": True, + } + } + runner._accumulated_valid_tokens = {"policy-a": 17} + return runner + + +def test_kill_session_rejects_missing_checkpoint_for_nonresident_lora_session(tmp_path): + runner = _build_runner(tmp_path, has_adapter=False) + + with pytest.raises(FileNotFoundError, match="no evicted checkpoint exists"): + runner.kill_session("policy-a", save_checkpoint=True) + + assert "policy-a" in runner._lora_session_specs + assert runner._accumulated_valid_tokens["policy-a"] == 17 + assert runner._adapter_manager.removed == [] + + +def test_kill_session_reuses_existing_evicted_checkpoint_for_nonresident_lora_session(tmp_path): + runner = _build_runner(tmp_path, has_adapter=False) + evicted_path = tmp_path / "adapters" / "evicted" / "policy-a" + evicted_path.mkdir(parents=True) + (evicted_path / "metadata.json").write_text('{"saved": true}', encoding="utf-8") + + result = runner.kill_session("policy-a", save_checkpoint=True) + promoted_path = tmp_path / "weights" / "policy-a" / "session_policy-a_final" + + assert result == { + "success": True, + "message": "LoRA session 'policy-a' killed successfully.", + "checkpoint_path": str(promoted_path), + } + assert (promoted_path / "metadata.json").read_text(encoding="utf-8") == '{"saved": true}' + assert "policy-a" not in runner._lora_session_specs + assert "policy-a" not in runner._accumulated_valid_tokens + assert runner._adapter_manager.removed == [] diff --git a/tests/server/runner/test_model_runner_session_registry.py b/tests/server/runner/test_model_runner_session_registry.py new file mode 100644 index 00000000..158ddd6d --- /dev/null +++ b/tests/server/runner/test_model_runner_session_registry.py @@ -0,0 +1,186 @@ +"""Tests for LoRA session-registry synchronization in ModelRunner.""" + +import importlib.util +from copy import deepcopy +from pathlib import Path + +import pytest +import torch + + +_MODULE_PATH = Path(__file__).resolve().parents[3] / "src" / "xorl" / "server" / "runner" / "model_runner.py" +_SPEC = importlib.util.spec_from_file_location("xorl_test_model_runner_session_registry", _MODULE_PATH) +assert _SPEC is not None and _SPEC.loader is not None +_MODULE = importlib.util.module_from_spec(_SPEC) +_SPEC.loader.exec_module(_MODULE) +ModelRunner = _MODULE.ModelRunner + + +pytestmark = [pytest.mark.cpu, pytest.mark.server] + + +def _session_spec(lr: float) -> dict: + return { + "base_model": "Qwen/Qwen3-8B", + "is_lora": True, + "lora_config": {"lora_rank": 4, "lora_alpha": 8}, + "optimizer_config": { + "type": "adamw", + "learning_rate": lr, + "weight_decay": 0.01, + "optimizer_dtype": "bf16", + "betas": [0.9, 0.95], + "eps": 1e-8, + "optimizer_kwargs": {}, + }, + } + + +class _FakeAdapterManager: + def __init__(self, lr: float) -> None: + self.session_specs = {"policy-a": _session_spec(lr)} + self.optim_step_calls = [] + + def has_adapter(self, model_id: str) -> bool: + return model_id in self.session_specs + + def get_adapter_session_spec(self, model_id: str) -> dict: + return deepcopy(self.session_specs[model_id]) + + def get_lr(self, model_id: str) -> float: + return self.session_specs[model_id]["optimizer_config"]["learning_rate"] + + def optim_step(self, model_id: str, lr: float, clip_value: float, *, accumulated_valid_tokens: int = 0) -> float: + self.optim_step_calls.append((model_id, lr, clip_value, accumulated_valid_tokens)) + self.session_specs[model_id]["optimizer_config"]["learning_rate"] = lr + return 7.5 + + def get_global_step(self, model_id: str) -> int: + return 3 + + +class _FakeCheckpointManager: + def __init__(self) -> None: + self.global_step = 11 + self.global_forward_backward_step = 13 + self.load_calls = [] + + def load_adapter_state(self, model_id, path=None, load_optimizer=True, lr=None): + self.load_calls.append((model_id, path, load_optimizer, lr)) + return {"success": True, "path": path, "model_id": model_id} + + +class _TinyModule(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.param = torch.nn.Parameter(torch.tensor([2.0])) + + +class _FakeOptimizer: + def __init__(self) -> None: + self.param_groups = [{"lr": 0.1}] + self.step_calls = 0 + self.zero_grad_calls = 0 + + def step(self) -> None: + self.step_calls += 1 + + def zero_grad(self, set_to_none=True) -> None: + self.zero_grad_calls += 1 + + +def _build_runner() -> ModelRunner: + runner = object.__new__(ModelRunner) + runner.rank = 0 + runner.is_sleeping = False + runner.lora_config = {"enable_lora": True, "merge_lora_interval": 0} + runner.train_config = {} + runner._accumulated_valid_tokens = {"policy-a": 11} + runner._lora_session_specs = {"policy-a": _session_spec(0.05)} + runner.global_step = 0 + runner.global_forward_backward_step = 0 + return runner + + +def test_optim_step_syncs_registered_lora_session_spec(monkeypatch): + runner = _build_runner() + runner._adapter_manager = _FakeAdapterManager(lr=0.05) + + monkeypatch.setattr(_MODULE, "synchronize", lambda: None) + + result = ModelRunner.optim_step(runner, gradient_clip=1.0, lr=0.25, model_id="policy-a") + + assert result["lr"] == pytest.approx(0.25) + assert runner._adapter_manager.optim_step_calls == [("policy-a", 0.25, 1.0, 11)] + assert runner._lora_session_specs["policy-a"]["optimizer_config"]["learning_rate"] == pytest.approx(0.25) + + +def test_load_adapter_state_syncs_registered_lora_session_spec(): + runner = _build_runner() + runner._adapter_manager = _FakeAdapterManager(lr=0.25) + runner._checkpoint_mgr = _FakeCheckpointManager() + + result = ModelRunner.load_adapter_state( + runner, + "policy-a", + path="/tmp/checkpoint", + load_optimizer=False, + lr=None, + ) + + assert result == { + "success": True, + "path": "/tmp/checkpoint", + "model_id": "policy-a", + } + assert runner._checkpoint_mgr.load_calls == [("policy-a", "/tmp/checkpoint", False, None)] + assert runner.global_step == 11 + assert runner.global_forward_backward_step == 13 + assert runner._lora_session_specs["policy-a"]["optimizer_config"]["learning_rate"] == pytest.approx(0.25) + + +def test_optim_step_preserves_distsignsgd_scaling_and_clip(monkeypatch): + runner = object.__new__(ModelRunner) + runner.rank = 0 + runner.is_sleeping = False + runner._adapter_manager = None + runner._use_distsignsgd = True + runner._accumulated_valid_tokens = {"default": 100} + runner._accumulated_active_microbatches = {"default": 2} + runner._accumulated_active_voter_total = {"default": 4} + runner.train_config = {"max_grad_norm": 1.0} + runner.lora_config = {"enable_lora": False, "merge_lora_interval": 0} + runner.model = _TinyModule() + runner.model.param.grad = torch.tensor([4.0]) + runner.optimizer = _FakeOptimizer() + runner.pp_enabled = False + runner.global_step = 0 + + captured = {} + + monkeypatch.setattr( + _MODULE, + "get_parallel_state", + lambda: type("ParallelState", (), {"fsdp_group": None, "pp_group": None})(), + ) + monkeypatch.setattr( + _MODULE, + "clip_gradients", + lambda model, clip_value, pp_enabled=False, pp_group=None: captured.update( + {"clip_value": clip_value, "grad": model.param.grad.item()} + ) + or 7.0, + ) + monkeypatch.setattr(_MODULE, "all_reduce", lambda value, group=None: value) + monkeypatch.setattr(_MODULE, "synchronize", lambda: None) + monkeypatch.setattr(_MODULE, "_maybe_merge_lora_util", lambda *args, **kwargs: None) + monkeypatch.setattr(_MODULE.torch.cuda, "empty_cache", lambda: None) + + result = ModelRunner.optim_step(runner, model_id="default") + + assert captured["clip_value"] == float("inf") + assert captured["grad"] == pytest.approx(1.0) + assert runner.optimizer.step_calls == 1 + assert runner.optimizer.zero_grad_calls == 1 + assert result["step"] == 1 + assert result["grad_norm"] == pytest.approx(7.0) diff --git a/tests/server/runner/test_opd_runner.py b/tests/server/runner/test_opd_runner.py new file mode 100644 index 00000000..9b9d3c31 --- /dev/null +++ b/tests/server/runner/test_opd_runner.py @@ -0,0 +1,475 @@ +from contextlib import nullcontext +from types import SimpleNamespace +from unittest.mock import Mock, patch + +import pytest +import torch +from safetensors.torch import load_file + +from tests._helpers.opd import make_teacher_files +from xorl.server.runner.model_runner import ModelRunner + + +pytestmark = [pytest.mark.cpu, pytest.mark.server] + + +def _make_opd_runner() -> ModelRunner: + runner = object.__new__(ModelRunner) + runner.rank = 0 + runner.world_size = 1 + runner.train_config = {} + runner.lm_head_fp32 = True + runner.pp_enabled = False + runner.model_fwd_context = nullcontext() + runner._opd_head_manager = None + runner._opd_head_config = None + runner._opd_hidden_cache = None + runner._opd_hidden_config = None + return runner + + +class _FakeTeacherOutput: + def __init__(self, last_hidden_state: torch.Tensor): + self.last_hidden_state = last_hidden_state + + +class _InputIdHiddenModel: + def __call__(self, input_ids, **_kwargs): + return _FakeTeacherOutput(input_ids.float().unsqueeze(-1)) + + +class _RecordingLmHead(torch.nn.Linear): + def __init__(self, hidden_size: int, vocab_size: int): + super().__init__(hidden_size, vocab_size, bias=False) + self.calls = 0 + self.last_input_shape = None + + def forward(self, input): + self.calls += 1 + self.last_input_shape = tuple(input.shape) + return super().forward(input) + + +class _NoopRoutingHandler: + def setup(self, *_args, **_kwargs): + return False + + +def test_forward_uses_no_grad_not_inference_mode(): + runner = object.__new__(ModelRunner) + runner.rank = 0 + runner.model = SimpleNamespace(config=SimpleNamespace(vocab_size=10)) + runner.pp_enabled = False + runner._adapter_manager = None + runner.global_forward_backward_step = 0 + runner._routing_handler = _NoopRoutingHandler() + runner._check_not_sleeping = lambda *_args, **_kwargs: None + runner._validate_single_tenant = lambda *_args, **_kwargs: None + + seen = {} + + def fake_forward_loop(*_args, **_kwargs): + seen["grad_enabled"] = torch.is_grad_enabled() + seen["inference_mode"] = torch.is_inference_mode_enabled() + return {"total_loss": 0.0, "global_valid_tokens": 1} + + runner._forward_loop = fake_forward_loop + + result = runner.forward([{"input_ids": torch.tensor([[1]])}], loss_fn="teacher_hidden_cache") + + assert result["step"] == 0 + assert seen == {"grad_enabled": False, "inference_mode": False} + + +@patch("xorl.server.runner.model_runner.get_parallel_state") +def test_opd_metrics_keep_opd_namespace(mock_parallel_state): + mock_parallel_state.return_value = Mock(dp_enabled=False, loss_parallel_enabled=False) + accumulated = {} + + ModelRunner._accumulate_loss_metrics( + accumulated, + { + "valid_tokens": 4, + "opd_kl": 0.5, + "opd_weighted_kl": 0.6, + "opd_num_teachers": 2, + "opd_profile_kl_compute_ms": 10.0, + }, + "opd_loss", + ) + ModelRunner._accumulate_loss_metrics( + accumulated, + { + "valid_tokens": 2, + "opd_kl": 0.2, + "opd_weighted_kl": 0.3, + "opd_num_teachers": 1, + "opd_profile_kl_compute_ms": 20.0, + }, + "opd_loss", + ) + + result = {} + ModelRunner._finalize_loss_metrics(accumulated, result) + + assert result["opd_kl"] == pytest.approx((0.5 * 4 + 0.2 * 2) / 6) + assert result["opd_weighted_kl"] == pytest.approx((0.6 * 4 + 0.3 * 2) / 6) + assert result["opd_num_teachers:max"] == 2 + assert result["opd_profile_kl_compute_ms"] == pytest.approx(30.0) + assert not any(key.startswith("is_opd") for key in result) + + +@patch("xorl.server.runner.model_runner.get_device_type", return_value="cpu") +@patch("xorl.server.runner.model_runner.get_parallel_state") +def test_opd_metrics_reduce_over_loss_group(mock_parallel_state, _mock_get_device_type): + loss_group = object() + mock_parallel_state.return_value = Mock(loss_parallel_enabled=True, loss_group=loss_group) + accumulated = {} + ModelRunner._accumulate_loss_metrics( + accumulated, + {"valid_tokens": 3, "opd_kl": 0.5, "opd_num_teachers": 2}, + "opd_loss", + ) + + groups = [] + + def fake_all_reduce(_tensor, op=None, group=None): + groups.append(group) + + with ( + patch("xorl.server.runner.model_runner.dist.is_available", return_value=True), + patch("xorl.server.runner.model_runner.dist.is_initialized", return_value=True), + patch("xorl.server.runner.model_runner.dist.all_reduce", side_effect=fake_all_reduce), + ): + result = {} + ModelRunner._finalize_loss_metrics(accumulated, result, "opd_loss") + + assert groups + assert all(group is loss_group for group in groups) + + +@patch("xorl.server.runner.model_runner.synchronize", lambda: None) +@patch("xorl.server.runner.model_runner.get_device_type", return_value="cpu") +@patch("xorl.server.runner.model_runner.get_parallel_state") +def test_forward_loop_accumulates_opd_metrics_without_metric_ops(mock_parallel_state, _mock_get_device_type): + mock_parallel_state.return_value = Mock( + cp_enabled=False, + loss_parallel_enabled=False, + dp_enabled=False, + ) + runner = object.__new__(ModelRunner) + runner.rank = 0 + runner.pp_enabled = False + runner.model_fwd_context = nullcontext() + runner._use_distsignsgd = False + runner._count_global_valid_tokens = lambda _micro_batches: torch.tensor(2) + runner._collect_per_token_outputs = Mock() + + def fake_compute_micro_batch_loss(_micro_batch, _loss_fn, _params): + loss = torch.tensor(4.0) + metrics = { + "valid_tokens": 2, + "opd_kl": 0.5, + "opd_weighted_kl": 0.75, + } + return loss, {}, metrics, None, SimpleNamespace() + + runner._compute_micro_batch_loss = fake_compute_micro_batch_loss + + result = runner._forward_loop( + [{"labels": torch.tensor([[1, 2]])}], + "opd_loss", + {}, + compute_backward=False, + ) + + assert result["total_loss"] == pytest.approx(2.0) + assert result["global_valid_tokens"] == 2 + assert result["opd_kl"] == pytest.approx(0.5) + assert result["opd_weighted_kl"] == pytest.approx(0.75) + + +@patch("xorl.server.runner.model_runner.get_device_type", return_value="cpu") +def test_opd_runner_masks_cache_indices_per_teacher(_mock_get_device_type, tmp_path): + torch.manual_seed(7) + vocab_size = 13 + hidden_size = 4 + seq_len = 4 + + teacher_heads = { + "0": torch.randn(vocab_size, hidden_size) / hidden_size**0.5, + "1": torch.randn(vocab_size, hidden_size) / hidden_size**0.5, + } + teacher_caches = { + "0": torch.randn(2, hidden_size) / hidden_size**0.5, + "1": torch.randn(12, hidden_size) / hidden_size**0.5, + } + teacher_files = make_teacher_files(tmp_path, teacher_heads, teacher_caches) + + runner = _make_opd_runner() + hidden_states = (torch.randn(1, seq_len, hidden_size) / hidden_size**0.5).requires_grad_(True) + student_weight = (torch.randn(vocab_size, hidden_size) / hidden_size**0.5).requires_grad_(True) + micro_batch = { + "labels": torch.tensor([[2, 3, 4, 5]]), + "teacher_ids": torch.tensor([[0, 0, 1, 1]]), + "teacher_cache_indices": torch.tensor([[0, 1, 10, 11]]), + "teacher_weights": torch.ones(1, seq_len), + } + params = { + "teacher_heads": teacher_files.heads, + "teacher_hidden_caches": teacher_files.hidden_caches, + "num_chunks": 2, + "opd_kl_backend": "streaming", + "opd_vocab_chunk_size": 5, + "opd_profile_timings": True, + } + + result = runner._compute_opd_micro_batch_loss( + hidden_states=hidden_states, + student_weight=student_weight, + micro_batch=micro_batch, + params=params, + ) + + assert result.loss.isfinite() + assert result.metrics["valid_tokens"] == seq_len + assert result.metrics["opd_num_teachers"] == 2 + assert result.metrics["opd_profile_hidden_fetch_ms"] >= 0.0 + assert result.metrics["opd_profile_head_prepare_ms"] >= 0.0 + assert result.metrics["opd_profile_kl_compute_ms"] >= 0.0 + assert result.metrics["opd_profile_total_ms"] >= result.metrics["opd_profile_kl_compute_ms"] + result.loss.backward() + assert hidden_states.grad is not None + assert student_weight.grad is not None + + +@patch("xorl.server.runner.model_runner.get_device_type", return_value="cpu") +def test_opd_runner_runs_lm_head_anchor_for_fsdp(_mock_get_device_type, tmp_path): + torch.manual_seed(11) + seq_len = 3 + hidden_size = 4 + teacher_hidden_size = 5 + vocab_size = 9 + teacher_heads = {"0": torch.randn(vocab_size, teacher_hidden_size) / teacher_hidden_size**0.5} + teacher_caches = {"0": torch.randn(seq_len, teacher_hidden_size) / teacher_hidden_size**0.5} + teacher_files = make_teacher_files(tmp_path, teacher_heads, teacher_caches) + + runner = _make_opd_runner() + hidden_states = (torch.randn(1, seq_len, hidden_size) / hidden_size**0.5).requires_grad_(True) + lm_head = _RecordingLmHead(hidden_size, vocab_size) + micro_batch = { + "labels": torch.tensor([[1, 2, 3]]), + "teacher_ids": torch.zeros(1, seq_len, dtype=torch.long), + "teacher_cache_indices": torch.arange(seq_len, dtype=torch.long).unsqueeze(0), + } + params = { + "teacher_heads": teacher_files.heads, + "teacher_hidden_caches": teacher_files.hidden_caches, + "opd_kl_backend": "streaming", + "opd_vocab_chunk_size": 4, + } + + result = runner._compute_opd_micro_batch_loss( + hidden_states=hidden_states, + student_weight=lm_head.weight, + micro_batch=micro_batch, + params=params, + student_lm_head=lm_head, + ) + + assert lm_head.calls == 1 + assert lm_head.last_input_shape == (1, hidden_size) + result.loss.backward() + assert hidden_states.grad is not None and hidden_states.grad.isfinite().all() + assert lm_head.weight.grad is not None and lm_head.weight.grad.isfinite().all() + + +def test_teacher_hidden_cache_splits_packed_batch_and_drops_padding(): + runner = _make_opd_runner() + hidden_states = torch.arange(1 * 8 * 2, dtype=torch.float32).reshape(1, 8, 2) + micro_batch = { + "num_samples": 2, + # Two real samples with lengths 3 and 2, then a padding segment that + # also starts at position 0. + "position_ids": torch.tensor([[0, 1, 2, 0, 1, 0, 1, 2]]), + } + + chunks = runner._teacher_hidden_chunks_from_batch(hidden_states, micro_batch) + + assert len(chunks) == 2 + assert torch.equal(chunks[0], hidden_states[0, 0:3]) + assert torch.equal(chunks[1], hidden_states[0, 3:5]) + + +def test_teacher_hidden_cache_contributor_skips_duplicate_cp_and_ep_ranks(): + runner = _make_opd_runner() + runner.rank = 3 + runner.world_size = 8 + + assert ( + runner._teacher_hidden_cache_contributor_key(SimpleNamespace(cp_enabled=True, cp_rank=1, ep_enabled=False)) + is None + ) + assert ( + runner._teacher_hidden_cache_contributor_key( + SimpleNamespace(cp_enabled=True, cp_rank=0, ep_enabled=False, dp_rank=2) + ) + == 2 + ) + assert ( + runner._teacher_hidden_cache_contributor_key(SimpleNamespace(cp_enabled=False, ep_enabled=True, ep_rank=1)) + is None + ) + + class FakeEpMesh: + @staticmethod + def get_local_rank(dim): + assert dim == "ep_fsdp" + return 3 + + assert ( + runner._teacher_hidden_cache_contributor_key( + SimpleNamespace(cp_enabled=False, ep_enabled=True, ep_rank=0, ep_fsdp_device_mesh=FakeEpMesh()) + ) + == 3 + ) + + +def test_teacher_hidden_cache_merge_preserves_logical_slice_order(): + chunks, indices = ModelRunner._merge_teacher_hidden_cache_payloads( + [ + {"rank": 4, "slice_key": 1, "chunks": [torch.ones(2, 2)]}, + None, + {"rank": 0, "slice_key": 0, "chunks": [torch.zeros(1, 2)]}, + ] + ) + + assert indices == [[0], [1, 2]] + assert torch.equal(torch.cat(chunks, dim=0), torch.tensor([[0.0, 0.0], [1.0, 1.0], [1.0, 1.0]])) + + +@patch("xorl.server.runner.model_runner.get_device_type", return_value="cpu") +@patch("xorl.server.runner.model_runner.gather_outputs") +@patch("xorl.server.runner.model_runner.get_parallel_state") +def test_teacher_hidden_cache_gathers_with_unified_sp_group(mock_parallel_state, mock_gather, _mock_device, tmp_path): + runner = _make_opd_runner() + runner.rank = 0 + runner.world_size = 1 + runner.model_fwd_context = nullcontext() + + class FakeModel: + def __call__(self, **_kwargs): + return SimpleNamespace(last_hidden_state=torch.arange(1 * 2 * 2, dtype=torch.float32).reshape(1, 2, 2)) + + runner.model = FakeModel() + mock_parallel_state.return_value = SimpleNamespace( + cp_enabled=True, + cp_size=4, + cp_rank=0, + sp_group="full-sp-group", + ep_enabled=False, + dp_rank=0, + ) + gathered = torch.arange(1 * 8 * 2, dtype=torch.float32).reshape(1, 8, 2) + mock_gather.return_value = gathered + + result = runner._forward_teacher_hidden_cache( + [ + { + "input_ids": torch.ones(1, 2, dtype=torch.long), + "_original_position_ids": torch.arange(8, dtype=torch.long).view(1, 8), + } + ], + {"teacher_hidden_cache_path": str(tmp_path / "teacher.safetensors")}, + ) + + mock_gather.assert_called_once() + assert mock_gather.call_args.kwargs["group"] == "full-sp-group" + assert mock_gather.call_args.kwargs["unpad_dim_size"] == 8 + assert result["teacher_hidden_cache"]["cache_indices_by_sample"] == [list(range(8))] + + +@patch("xorl.server.runner.model_runner.get_device_type", return_value="cpu") +@patch("xorl.server.runner.model_runner.get_parallel_state") +def test_teacher_hidden_cache_trims_with_gathered_sp_labels(mock_parallel_state, _mock_get_device_type, tmp_path): + runner = _make_opd_runner() + runner.model = _InputIdHiddenModel() + mock_parallel_state.return_value = Mock( + cp_enabled=True, + cp_rank=0, + cp_size=2, + ulysses_group=object(), + ep_enabled=False, + dp_rank=0, + ) + + full_hidden = torch.arange(6, dtype=torch.float32).reshape(1, 6, 1) + full_labels = torch.tensor([[-100, -100, -100, -100, 9, -100]]) + + def fake_gather_outputs(tensor, **_kwargs): + return full_hidden if torch.is_floating_point(tensor) else full_labels + + cache_path = tmp_path / "teacher_hidden.safetensors" + with patch("xorl.server.runner.model_runner.gather_outputs", side_effect=fake_gather_outputs): + result = runner._forward_teacher_hidden_cache( + [ + { + "input_ids": torch.tensor([[1, 2, 3]]), + "labels": torch.tensor([[-100, -100, -100]]), + "_original_position_ids": torch.arange(6, dtype=torch.long).unsqueeze(0), + } + ], + { + "teacher_hidden_cache_path": str(cache_path), + "teacher_hidden_cache_dtype": "float32", + }, + ) + + saved = load_file(str(cache_path))["hidden_states"] + assert torch.equal(saved, full_hidden[0, :5]) + assert result["teacher_hidden_cache"]["cache_indices_by_sample"] == [list(range(5))] + + +@patch("xorl.server.runner.model_runner.get_device_type", return_value="cpu") +@patch("xorl.server.runner.model_runner.get_parallel_state") +def test_teacher_hidden_cache_writer_gathers_all_batch_ranks( + mock_parallel_state, + _mock_get_device_type, + tmp_path, +): + runner = _make_opd_runner() + runner.world_size = 2 + runner.model = _InputIdHiddenModel() + mock_parallel_state.return_value = Mock(cp_enabled=False, ep_enabled=False, dp_rank=0) + + remote_chunk = torch.tensor([[10.0], [11.0], [12.0]]) + + def fake_gather_object(payload, object_gather_list, dst): + assert dst == 0 + object_gather_list[:] = [payload, {"rank": 1, "slice_key": 1, "chunks": [remote_chunk]}] + + cache_path = tmp_path / "teacher_hidden.safetensors" + with ( + patch("xorl.server.runner.model_runner.dist.is_available", return_value=True), + patch("xorl.server.runner.model_runner.dist.is_initialized", return_value=True), + patch("xorl.server.runner.model_runner.dist.get_world_size", return_value=2), + patch("xorl.server.runner.model_runner.dist.gather_object", side_effect=fake_gather_object), + patch("xorl.server.runner.model_runner.dist.broadcast_object_list"), + ): + result = runner._forward_teacher_hidden_cache( + [ + { + "input_ids": torch.tensor([[1, 2, 3]]), + "labels": torch.tensor([[1, 2, -100]]), + } + ], + { + "teacher_hidden_cache_path": str(cache_path), + "teacher_hidden_cache_dtype": "float32", + }, + ) + + saved = load_file(str(cache_path))["hidden_states"] + assert torch.equal(saved, torch.tensor([[1.0], [2.0], [10.0], [11.0], [12.0]])) + assert result["teacher_hidden_cache"]["num_tokens"] == 5 + assert result["teacher_hidden_cache"]["cache_indices_by_sample"] == [[0, 1], [2, 3, 4]] diff --git a/tests/server/runner/test_routing_replay_handler.py b/tests/server/runner/test_routing_replay_handler.py new file mode 100644 index 00000000..4fcb40dc --- /dev/null +++ b/tests/server/runner/test_routing_replay_handler.py @@ -0,0 +1,139 @@ +from types import SimpleNamespace + +import torch + +from xorl.server.runner.utils import routing_replay_handler as rrh + + +def _routing(start: int, length: int) -> list[list[list[int]]]: + return [[[start + i, start + i + 1000]] for i in range(length)] + + +def _handler() -> rrh.RoutingReplayHandler: + return rrh.RoutingReplayHandler(torch.nn.Module()) + + +def test_sp_routing_uses_actual_position_ids_length_without_128_padding(monkeypatch): + monkeypatch.setattr(rrh, "get_parallel_state", lambda: SimpleNamespace(cp_enabled=True, cp_size=4, cp_rank=1)) + micro_batches = [ + { + "input_ids": torch.zeros(1, 93, dtype=torch.long), + "position_ids": torch.arange(372, dtype=torch.long).view(1, 372), + "num_samples": 1, + } + ] + + per_mb = _handler()._build_per_mb_routing(micro_batches, [_routing(0, 372)], num_layers_in_data=1, topk=2) + + assert len(per_mb) == 1 + assert per_mb[0].shape == (93, 1, 2) + assert per_mb[0][0, 0, 0].item() == 93 + assert per_mb[0][-1, 0, 0].item() == 185 + + +def test_sp_routing_pads_to_actual_position_ids_length(monkeypatch): + monkeypatch.setattr(rrh, "get_parallel_state", lambda: SimpleNamespace(cp_enabled=True, cp_size=4, cp_rank=3)) + micro_batches = [ + { + "input_ids": torch.zeros(1, 96, dtype=torch.long), + "position_ids": torch.arange(384, dtype=torch.long).view(1, 384), + "num_samples": 1, + } + ] + + per_mb = _handler()._build_per_mb_routing(micro_batches, [_routing(0, 372)], num_layers_in_data=1, topk=2) + + assert per_mb[0].shape == (96, 1, 2) + assert per_mb[0][0, 0, 0].item() == 288 + assert per_mb[0][83, 0, 0].item() == 371 + assert per_mb[0][84, 0].tolist() == [0, 1] + + +def test_sp_routing_slices_unpacked_rows_independently(monkeypatch): + monkeypatch.setattr(rrh, "get_parallel_state", lambda: SimpleNamespace(cp_enabled=True, cp_size=4, cp_rank=1)) + micro_batches = [ + { + "input_ids": torch.zeros(2, 3, dtype=torch.long), + "position_ids": torch.arange(24, dtype=torch.long).view(2, 12), + "num_samples": 2, + } + ] + + per_mb = _handler()._build_per_mb_routing( + micro_batches, + [_routing(0, 12), _routing(100, 12)], + num_layers_in_data=1, + topk=2, + ) + + assert per_mb[0].shape == (6, 1, 2) + assert per_mb[0][:, 0, 0].tolist() == [3, 4, 5, 103, 104, 105] + + +def test_ringattn_routing_uses_zigzag_layout_before_sp_slice(monkeypatch): + expected_by_rank = { + 0: [0, 1], + 1: [6, 7], + 2: [2, 3], + 3: [4, 5], + } + for cp_rank, expected in expected_by_rank.items(): + monkeypatch.setattr( + rrh, + "get_parallel_state", + lambda cp_rank=cp_rank: SimpleNamespace( + cp_enabled=True, + cp_size=4, + cp_rank=cp_rank, + ringattn_size=2, + ), + ) + micro_batches = [ + { + "input_ids": torch.zeros(1, 2, dtype=torch.long), + "position_ids": torch.tensor([[0, 1, 6, 7, 2, 3, 4, 5]], dtype=torch.long), + "_original_position_ids": torch.arange(8, dtype=torch.long).view(1, 8), + "num_samples": 1, + } + ] + + per_mb = _handler()._build_per_mb_routing(micro_batches, [_routing(0, 8)], num_layers_in_data=1, topk=2) + + assert per_mb[0].shape == (2, 1, 2) + assert per_mb[0][:, 0, 0].tolist() == expected + + +def test_ringattn_routing_zigzag_respects_packed_document_boundaries(monkeypatch): + monkeypatch.setattr( + rrh, + "get_parallel_state", + lambda: SimpleNamespace(cp_enabled=True, cp_size=2, cp_rank=0, ringattn_size=2), + ) + micro_batches = [ + { + "input_ids": torch.zeros(1, 4, dtype=torch.long), + "position_ids": torch.tensor([[0, 3, 0, 3, 1, 2, 1, 2]], dtype=torch.long), + "_original_position_ids": torch.tensor([[0, 1, 2, 3, 0, 1, 2, 3]], dtype=torch.long), + "num_samples": 2, + } + ] + + per_mb = _handler()._build_per_mb_routing( + micro_batches, + [_routing(0, 4), _routing(4, 4)], + num_layers_in_data=1, + topk=2, + ) + + assert per_mb[0].shape == (4, 1, 2) + assert per_mb[0][:, 0, 0].tolist() == [0, 3, 4, 7] + + +def test_routing_truncates_excess_to_micro_batch_size(monkeypatch): + monkeypatch.setattr(rrh, "get_parallel_state", lambda: SimpleNamespace(cp_enabled=False)) + micro_batches = [{"input_ids": torch.zeros(1, 3, dtype=torch.long), "num_samples": 1}] + + per_mb = _handler()._build_per_mb_routing(micro_batches, [_routing(0, 4)], num_layers_in_data=1, topk=2) + + assert per_mb[0].shape == (3, 1, 2) + assert per_mb[0][:, 0, 0].tolist() == [0, 1, 2] diff --git a/tests/server/runner/test_runner_dispatcher.py b/tests/server/runner/test_runner_dispatcher.py new file mode 100644 index 00000000..dc2a2091 --- /dev/null +++ b/tests/server/runner/test_runner_dispatcher.py @@ -0,0 +1,106 @@ +from types import SimpleNamespace + +import pytest +import torch + +import xorl.server.runner.runner_dispatcher as runner_dispatcher_module +from xorl.server.runner.runner_dispatcher import RunnerDispatcher + + +pytestmark = [pytest.mark.cpu, pytest.mark.server] + + +class _FakeEPMesh: + def __init__(self, ep_fsdp_rank: int) -> None: + self.ep_fsdp_rank = ep_fsdp_rank + + def get_local_rank(self, name: str) -> int: + assert name == "ep_fsdp" + return self.ep_fsdp_rank + + +def _dispatcher(rank: int, world_size: int) -> RunnerDispatcher: + dispatcher = object.__new__(RunnerDispatcher) + dispatcher.rank = rank + dispatcher.world_size = world_size + return dispatcher + + +def _batch(batch_id: int, *, num_samples: int = 1) -> dict: + return { + "input_ids": [[batch_id, batch_id + 1]], + "labels": [[batch_id + 2, batch_id + 3]], + "position_ids": [[0, 1]], + "num_samples": num_samples, + } + + +def _parallel_state(**overrides): + return SimpleNamespace( + cp_size=overrides.get("cp_size", 1), + pp_enabled=overrides.get("pp_enabled", False), + pp_size=overrides.get("pp_size", 1), + ep_enabled=overrides.get("ep_enabled", False), + ep_size=overrides.get("ep_size", 1), + dp_shard_in_ep_size=overrides.get("dp_shard_in_ep_size", 1), + ep_fsdp_device_mesh=overrides.get("ep_fsdp_device_mesh"), + ) + + +def test_select_batches_keeps_existing_dp_distribution_without_ep(monkeypatch): + monkeypatch.setattr(runner_dispatcher_module, "get_parallel_state", lambda: _parallel_state()) + + batches = [_batch(10), _batch(20), _batch(30), _batch(40)] + my_batches, routed_experts, routed_logits = _dispatcher(rank=2, world_size=4)._select_and_prepare_batches( + batches, + routed_experts=["r0", "r1", "r2", "r3"], + routed_expert_logits=["l0", "l1", "l2", "l3"], + ) + + assert len(my_batches) == 1 + assert torch.equal(my_batches[0]["input_ids"], torch.tensor([[30, 31]])) + assert routed_experts == ["r2"] + assert routed_logits == ["l2"] + + +def test_select_batches_broadcasts_one_slice_to_all_ep_ranks(monkeypatch): + state = _parallel_state( + ep_enabled=True, + ep_size=8, + dp_shard_in_ep_size=1, + ep_fsdp_device_mesh=_FakeEPMesh(ep_fsdp_rank=0), + ) + monkeypatch.setattr(runner_dispatcher_module, "get_parallel_state", lambda: state) + + my_batches, routed_experts, routed_logits = _dispatcher(rank=5, world_size=8)._select_and_prepare_batches( + [_batch(10, num_samples=2)], + routed_experts=["r0", "r1"], + routed_expert_logits=["l0", "l1"], + ) + + assert len(my_batches) == 1 + assert torch.equal(my_batches[0]["input_ids"], torch.tensor([[10, 11]])) + assert my_batches[0]["num_samples"] == 2 + assert routed_experts == ["r0", "r1"] + assert routed_logits == ["l0", "l1"] + + +def test_select_batches_uses_ep_fsdp_rank_for_distinct_ep_batch_slices(monkeypatch): + state = _parallel_state( + ep_enabled=True, + ep_size=4, + dp_shard_in_ep_size=2, + ep_fsdp_device_mesh=_FakeEPMesh(ep_fsdp_rank=1), + ) + monkeypatch.setattr(runner_dispatcher_module, "get_parallel_state", lambda: state) + + my_batches, routed_experts, routed_logits = _dispatcher(rank=6, world_size=8)._select_and_prepare_batches( + [_batch(10), _batch(20)], + routed_experts=["r0", "r1"], + routed_expert_logits=["l0", "l1"], + ) + + assert len(my_batches) == 1 + assert torch.equal(my_batches[0]["input_ids"], torch.tensor([[20, 21]])) + assert routed_experts == ["r1"] + assert routed_logits == ["l1"] diff --git a/tests/server/runner/test_runner_dispatcher_forward.py b/tests/server/runner/test_runner_dispatcher_forward.py new file mode 100644 index 00000000..e300c627 --- /dev/null +++ b/tests/server/runner/test_runner_dispatcher_forward.py @@ -0,0 +1,79 @@ +"""Tests for forward-path model_id handling in RunnerDispatcher.""" + +import asyncio +import importlib.util +from pathlib import Path + +import pytest + +from xorl.server.protocol.operations import ModelPassData + + +_MODULE_PATH = Path(__file__).resolve().parents[3] / "src" / "xorl" / "server" / "runner" / "runner_dispatcher.py" +_SPEC = importlib.util.spec_from_file_location("xorl_test_runner_dispatcher", _MODULE_PATH) +assert _SPEC is not None and _SPEC.loader is not None +_MODULE = importlib.util.module_from_spec(_SPEC) +_SPEC.loader.exec_module(_MODULE) +RunnerDispatcher = _MODULE.RunnerDispatcher + + +pytestmark = [pytest.mark.cpu, pytest.mark.server] + + +class _FakeAdapterCoordinator: + def __init__(self): + self.calls = [] + + def auto_load_if_evicted(self, model_id: str): + self.calls.append(model_id) + return True, "/tmp/adapter-state" + + +def test_handle_compute_rank0_scatter_uses_model_id_for_forward(monkeypatch): + dispatcher = object.__new__(RunnerDispatcher) + dispatcher.rank = 0 + dispatcher._adapter_coordinator = _FakeAdapterCoordinator() + + captured = {} + + def fake_select_and_prepare_batches(batches, routed_experts=None, routed_expert_logits=None): + return batches, routed_experts, routed_expert_logits + + def fake_execute_and_gather( + my_batches, + loss_fn, + loss_fn_params, + routed_experts, + cp_enabled, + parallel_state, + *, + with_backward, + model_id, + is_rank0, + routed_expert_logits=None, + ): + captured["model_id"] = model_id + captured["with_backward"] = with_backward + captured["is_rank0"] = is_rank0 + return {"total_loss": 0.5} + + dispatcher._select_and_prepare_batches = fake_select_and_prepare_batches + dispatcher._execute_and_gather = fake_execute_and_gather + + monkeypatch.setattr(_MODULE, "get_parallel_state", lambda: type("PS", (), {"cp_enabled": False})()) + + payload = ModelPassData( + batches=[{"model_input": {"input_ids": [1]}, "loss_fn_inputs": {"labels": [1]}}], + model_id="adapter-42", + ) + + result = asyncio.run(dispatcher._handle_compute_rank0_scatter({"payload": payload}, with_backward=False)) + + assert dispatcher._adapter_coordinator.calls == ["adapter-42"] + assert captured == { + "model_id": "adapter-42", + "with_backward": False, + "is_rank0": True, + } + assert result["auto_loaded"] is True + assert result["auto_load_path"] == "/tmp/adapter-state" diff --git a/tests/server/runner/test_runner_dispatcher_load_state.py b/tests/server/runner/test_runner_dispatcher_load_state.py new file mode 100644 index 00000000..0a355160 --- /dev/null +++ b/tests/server/runner/test_runner_dispatcher_load_state.py @@ -0,0 +1,110 @@ +"""Tests for load_state delegation in RunnerDispatcher.""" + +import asyncio +import importlib.util +from pathlib import Path + +import pytest + +from xorl.server.protocol.operations import AdapterStateData, LoadStateData + + +_MODULE_PATH = Path(__file__).resolve().parents[3] / "src" / "xorl" / "server" / "runner" / "runner_dispatcher.py" +_SPEC = importlib.util.spec_from_file_location("xorl_test_runner_dispatcher_load_state", _MODULE_PATH) +assert _SPEC is not None and _SPEC.loader is not None +_MODULE = importlib.util.module_from_spec(_SPEC) +_SPEC.loader.exec_module(_MODULE) +RunnerDispatcher = _MODULE.RunnerDispatcher + + +pytestmark = [pytest.mark.cpu, pytest.mark.server] + + +class _FakeAdapterCoordinator: + def __init__(self): + self.calls = [] + + async def handle_load_adapter_state(self, command_dict): + self.calls.append(command_dict) + return {"success": True, "model_id": command_dict["payload"].model_id, "step": 7} + + +class _FakeTrainerMultiAdapter: + def __init__(self): + self.adapter_manager = object() + self.step = 99 + self.load_state_calls = [] + + def load_state(self, checkpoint_path, load_optimizer=True, model_id=None): + self.load_state_calls.append((checkpoint_path, load_optimizer, model_id)) + return {"success": True} + + +class _FakeTrainerSingleTenant: + def __init__(self): + self.adapter_manager = None + self.step = 99 + self.load_state_calls = [] + + def load_state(self, checkpoint_path, load_optimizer=True, model_id=None): + self.load_state_calls.append((checkpoint_path, load_optimizer, model_id)) + return {"success": True, "model_id": model_id} + + +def test_handle_load_state_uses_adapter_coordinator_for_multi_adapter(tmp_path): + checkpoint_path = tmp_path / "checkpoint" + checkpoint_path.mkdir() + + dispatcher = object.__new__(RunnerDispatcher) + dispatcher.rank = 0 + dispatcher.trainer = _FakeTrainerMultiAdapter() + dispatcher._adapter_coordinator = _FakeAdapterCoordinator() + + result = asyncio.run( + dispatcher._handle_load_state( + { + "payload": LoadStateData( + checkpoint_path=str(checkpoint_path), + load_optimizer=False, + model_id="policy-a", + ) + } + ) + ) + + assert result == {"success": True, "model_id": "policy-a", "step": 7} + assert dispatcher.trainer.step == 0 + assert dispatcher.trainer.load_state_calls == [] + assert len(dispatcher._adapter_coordinator.calls) == 1 + payload = dispatcher._adapter_coordinator.calls[0]["payload"] + assert isinstance(payload, AdapterStateData) + assert payload.model_id == "policy-a" + assert payload.path == str(checkpoint_path) + assert payload.load_optimizer is False + + +def test_handle_load_state_uses_trainer_load_state_without_adapter_manager(tmp_path): + checkpoint_path = tmp_path / "checkpoint" + checkpoint_path.mkdir() + + dispatcher = object.__new__(RunnerDispatcher) + dispatcher.rank = 0 + dispatcher.trainer = _FakeTrainerSingleTenant() + dispatcher._adapter_coordinator = _FakeAdapterCoordinator() + + result = asyncio.run( + dispatcher._handle_load_state( + { + "payload": LoadStateData( + checkpoint_path=str(checkpoint_path), + load_optimizer=True, + model_id="default", + ) + } + ) + ) + + assert result == {"success": True, "model_id": "default"} + assert dispatcher.trainer.step == 0 + assert dispatcher.trainer.load_state_calls == [(str(checkpoint_path), True, "default")] + assert dispatcher._adapter_coordinator.calls == [] diff --git a/tests/server/runner/test_runner_dispatcher_session_ops.py b/tests/server/runner/test_runner_dispatcher_session_ops.py new file mode 100644 index 00000000..0047b6fb --- /dev/null +++ b/tests/server/runner/test_runner_dispatcher_session_ops.py @@ -0,0 +1,164 @@ +"""Tests for save/session operation handling in RunnerDispatcher.""" + +import asyncio +import importlib.util +from pathlib import Path + +import pytest + +from xorl.server.protocol.operations import RegisterSessionData, SaveLoraOnlyData, SaveStateData +from xorl.server.protocol.orchestrator_runner import RunnerDispatchCommand + + +_MODULE_PATH = Path(__file__).resolve().parents[3] / "src" / "xorl" / "server" / "runner" / "runner_dispatcher.py" +_SPEC = importlib.util.spec_from_file_location("xorl_test_runner_dispatcher_session_ops", _MODULE_PATH) +assert _SPEC is not None and _SPEC.loader is not None +_MODULE = importlib.util.module_from_spec(_SPEC) +_SPEC.loader.exec_module(_MODULE) +RunnerDispatcher = _MODULE.RunnerDispatcher + + +pytestmark = [pytest.mark.cpu, pytest.mark.server] + + +class _FakeAdapterCoordinator: + def __init__(self): + self.auto_load_calls = [] + self.register_session_calls = [] + + def auto_load_if_evicted(self, model_id: str, *, allow_fresh_materialization: bool = True): + self.auto_load_calls.append( + { + "model_id": model_id, + "allow_fresh_materialization": allow_fresh_materialization, + } + ) + return False, None + + async def handle_register_session(self, command_dict): + self.register_session_calls.append(command_dict) + payload = command_dict["payload"] + return {"registered": True, "model_id": payload.model_id} + + +class _FakeTrainer: + def __init__(self): + self.adapter_manager = object() + self.lora_config = {"enable_lora": True} + self.save_state_calls = [] + self.save_lora_only_calls = [] + + def save_state(self, checkpoint_path, save_optimizer=True, model_id=None): + self.save_state_calls.append((checkpoint_path, save_optimizer, model_id)) + return {"success": True} + + def save_lora_only(self, lora_path, model_id="default"): + self.save_lora_only_calls.append((lora_path, model_id)) + return {"success": True} + + +def test_handle_save_state_requires_real_checkpoint_for_nonresident_adapter(tmp_path): + dispatcher = object.__new__(RunnerDispatcher) + dispatcher.rank = 0 + dispatcher.trainer = _FakeTrainer() + dispatcher._adapter_coordinator = _FakeAdapterCoordinator() + + result = asyncio.run( + dispatcher._handle_save_state( + { + "payload": SaveStateData( + checkpoint_path=str(tmp_path / "checkpoint"), + save_optimizer=True, + model_id="policy-a", + ) + } + ) + ) + + assert result["checkpoint_path"] == str(tmp_path / "checkpoint") + assert dispatcher._adapter_coordinator.auto_load_calls == [ + { + "model_id": "policy-a", + "allow_fresh_materialization": False, + } + ] + assert dispatcher.trainer.save_state_calls == [(str(tmp_path / "checkpoint"), True, "policy-a")] + + +def test_handle_save_lora_only_requires_real_checkpoint_for_nonresident_adapter(tmp_path): + dispatcher = object.__new__(RunnerDispatcher) + dispatcher.rank = 0 + dispatcher.trainer = _FakeTrainer() + dispatcher._adapter_coordinator = _FakeAdapterCoordinator() + + result = asyncio.run( + dispatcher._handle_save_lora_only( + { + "payload": SaveLoraOnlyData( + lora_path=str(tmp_path / "adapter"), + model_id="policy-b", + ) + } + ) + ) + + assert result["lora_path"] == str(tmp_path / "adapter") + assert dispatcher._adapter_coordinator.auto_load_calls == [ + { + "model_id": "policy-b", + "allow_fresh_materialization": False, + } + ] + assert dispatcher.trainer.save_lora_only_calls == [(str(tmp_path / "adapter"), "policy-b")] + + +def test_handle_register_session_delegates_to_adapter_coordinator(): + dispatcher = object.__new__(RunnerDispatcher) + dispatcher.rank = 0 + dispatcher._adapter_coordinator = _FakeAdapterCoordinator() + payload = RegisterSessionData( + model_id="policy-c", + session_spec={ + "base_model": "Qwen/Qwen3-8B", + "is_lora": True, + "lora_config": {"lora_rank": 8, "lora_alpha": 16}, + "optimizer_config": {"type": "adamw", "learning_rate": 1e-4}, + }, + materialize=True, + ) + + result = asyncio.run(dispatcher._handle_register_session({"payload": payload})) + + assert result == {"registered": True, "model_id": "policy-c"} + assert dispatcher._adapter_coordinator.register_session_calls == [{"payload": payload}] + + +def test_handle_request_rank0_fails_register_session_on_cross_rank_worker_error(monkeypatch): + dispatcher = object.__new__(RunnerDispatcher) + dispatcher.rank = 0 + dispatcher.world_size = 2 + dispatcher.cpu_group = object() + dispatcher._worker_error = None + + async def _handle_register_session(command_dict): + return {"registered": True, "model_id": command_dict["payload"].model_id} + + dispatcher._handle_register_session = _handle_register_session + dispatcher._sync_error_state = lambda: "rank 1: Session registration failed: boom" + + monkeypatch.setattr(_MODULE.dist, "broadcast_object_list", lambda *args, **kwargs: None) + + request = RunnerDispatchCommand.create( + "register_session", + RegisterSessionData( + model_id="policy-c", + session_spec={"base_model": "Qwen/Qwen3-8B", "is_lora": True}, + materialize=False, + ), + request_id="req-register-session", + ) + + response = asyncio.run(dispatcher._handle_request_rank0(request)) + + assert response.success is False + assert response.error == "Cross-rank error: rank 1: Session registration failed: boom" diff --git a/tests/server/test_server_arguments.py b/tests/server/test_server_arguments.py index 1f63a119..c33491d0 100644 --- a/tests/server/test_server_arguments.py +++ b/tests/server/test_server_arguments.py @@ -1,12 +1,63 @@ +import importlib.util +import sys +import types +from pathlib import Path +from unittest.mock import patch + import pytest +import torch import yaml -from xorl.server.launcher import load_server_arguments - pytestmark = [pytest.mark.cpu, pytest.mark.server] +def _load_server_arguments_fn(): + module_path = Path(__file__).resolve().parents[2] / "src" / "xorl" / "server" / "launcher.py" + spec = importlib.util.spec_from_file_location("xorl_test_launcher", module_path) + assert spec is not None and spec.loader is not None + + fake_api_server_pkg = types.ModuleType("xorl.server.api_server") + fake_api_server_pkg.__path__ = [] + fake_api_server_mod = types.ModuleType("xorl.server.api_server.server") + fake_api_server_mod.APIServer = object + fake_api_server_pkg.server = fake_api_server_mod + + fake_orchestrator_pkg = types.ModuleType("xorl.server.orchestrator") + fake_orchestrator_pkg.__path__ = [] + fake_orchestrator_mod = types.ModuleType("xorl.server.orchestrator.orchestrator") + fake_orchestrator_mod.Orchestrator = object + fake_orchestrator_pkg.orchestrator = fake_orchestrator_mod + + fake_utils_pkg = types.ModuleType("xorl.server.utils") + fake_utils_pkg.__path__ = [] + fake_network_mod = types.ModuleType("xorl.server.utils.network") + fake_network_mod.read_address_file = lambda *args, **kwargs: None + fake_utils_pkg.network = fake_network_mod + fake_session_spec_mod = types.ModuleType("xorl.server.session_spec") + fake_session_spec_mod.build_default_session_spec = lambda *args, **kwargs: None + + module = importlib.util.module_from_spec(spec) + with patch.dict( + sys.modules, + { + "xorl.server.api_server": fake_api_server_pkg, + "xorl.server.api_server.server": fake_api_server_mod, + "xorl.server.orchestrator": fake_orchestrator_pkg, + "xorl.server.orchestrator.orchestrator": fake_orchestrator_mod, + "xorl.server.session_spec": fake_session_spec_mod, + "xorl.server.utils": fake_utils_pkg, + "xorl.server.utils.network": fake_network_mod, + }, + ): + spec.loader.exec_module(module) + + return module.load_server_arguments + + +load_server_arguments = _load_server_arguments_fn() + + def test_load_server_arguments_threads_signsgd_through_nested_config(tmp_path): config_path = tmp_path / "server_config.yaml" config_path.write_text( @@ -27,7 +78,144 @@ def test_load_server_arguments_threads_signsgd_through_nested_config(tmp_path): args = load_server_arguments(str(config_path)) assert args.optimizer == "signsgd" + assert args.load_weights_mode == "grouped" assert args.to_config_dict()["train"]["optimizer"] == "signsgd" + assert args.to_config_dict()["train"]["load_weights_mode"] == "grouped" + + +def test_load_server_arguments_threads_distsignsgd_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": "distsignsgd", + "output_dir": str(tmp_path / "outputs"), + }, + } + ), + encoding="utf-8", + ) + + args = load_server_arguments(str(config_path)) + + assert args.optimizer == "distsignsgd" + assert args.to_config_dict()["train"]["optimizer"] == "distsignsgd" + + +def test_load_server_arguments_threads_adapter_state_load_mode_into_lora_config(tmp_path): + config_path = tmp_path / "server_config.yaml" + config_path.write_text( + yaml.safe_dump( + { + "model": { + "model_path": "Qwen/Qwen3-8B", + }, + "lora": { + "enable_lora": True, + "adapter_state_load_mode": "rank0_broadcast", + }, + } + ), + encoding="utf-8", + ) + + args = load_server_arguments(str(config_path)) + + assert args.adapter_state_load_mode == "rank0_broadcast" + assert args.to_config_dict()["lora"]["adapter_state_load_mode"] == "rank0_broadcast" + + +def test_load_server_arguments_threads_activation_gpu_limit_into_train_config(tmp_path): + config_path = tmp_path / "server_config.yaml" + config_path.write_text( + yaml.safe_dump( + { + "model": { + "model_path": "Qwen/Qwen3-8B", + }, + "train": { + "enable_activation_offload": True, + "activation_gpu_limit": 0.25, + }, + } + ), + encoding="utf-8", + ) + + args = load_server_arguments(str(config_path)) + + assert args.activation_gpu_limit == pytest.approx(0.25) + assert args.to_config_dict()["train"]["activation_gpu_limit"] == pytest.approx(0.25) + + +def test_load_server_arguments_rejects_broadcast_load_weights_mode(tmp_path): + config_path = tmp_path / "server_config.yaml" + config_path.write_text( + yaml.safe_dump( + { + "model": { + "model_path": "Qwen/Qwen3-8B", + }, + "train": { + "load_weights_mode": "broadcast", + "output_dir": str(tmp_path / "outputs"), + }, + } + ), + encoding="utf-8", + ) + + with pytest.raises(ValueError, match="Unsupported load_weights_mode"): + load_server_arguments(str(config_path)) + + +def test_load_server_arguments_rejects_merge_lora_interval_for_server_multi_adapter(tmp_path): + config_path = tmp_path / "server_config.yaml" + config_path.write_text( + yaml.safe_dump( + { + "model": { + "model_path": "Qwen/Qwen3-8B", + }, + "lora": { + "enable_lora": True, + "merge_lora_interval": 16, + }, + } + ), + encoding="utf-8", + ) + + with pytest.raises(ValueError, match="merge_lora_interval is not supported"): + load_server_arguments(str(config_path)) + + +def test_load_server_arguments_rejects_pipeline_parallel_multi_adapter_lora(tmp_path): + config_path = tmp_path / "server_config.yaml" + config_path.write_text( + yaml.safe_dump( + { + "model": { + "model_path": "Qwen/Qwen3-8B", + }, + "train": { + "pipeline_parallel_size": 2, + }, + "lora": { + "enable_lora": True, + "adapter_state_load_mode": "rank0_broadcast", + }, + } + ), + encoding="utf-8", + ) + + with pytest.raises(ValueError, match="pipeline_parallel_size > 1 is not supported"): + load_server_arguments(str(config_path)) def test_load_server_arguments_threads_muon_gram_newton_schulz_through_nested_config(tmp_path): @@ -40,10 +228,13 @@ def test_load_server_arguments_threads_muon_gram_newton_schulz_through_nested_co }, "train": { "optimizer": "muon", + "optimizer_dtype": "bf16", + "muon_lr": 0.03, "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_grouped_gram_ns_fp32_byte_limit": 23, "muon_grad_dtype": "fp32", "muon_update_dtype": "bf16", "muon_force_momentum_path": True, @@ -55,17 +246,65 @@ def test_load_server_arguments_threads_muon_gram_newton_schulz_through_nested_co ) args = load_server_arguments(str(config_path)) - train_config = args.to_config_dict()["train"] + optimizer_kwargs = args.to_config_dict()["train"]["optimizer_kwargs"] + assert optimizer_kwargs["muon_lr"] == pytest.approx(0.03) + assert optimizer_kwargs["muon_ns_algorithm"] == "gram_newton_schulz" + assert optimizer_kwargs["muon_ns_use_quack_kernels"] is False + assert optimizer_kwargs["muon_gram_ns_num_restarts"] == 2 + assert optimizer_kwargs["muon_gram_ns_restart_iterations"] == [2] + assert optimizer_kwargs["muon_momentum_dtype"] == torch.bfloat16 + assert optimizer_kwargs["muon_grad_dtype"] == torch.float32 + assert optimizer_kwargs["muon_update_dtype"] == torch.bfloat16 + assert optimizer_kwargs["muon_force_momentum_path"] is True + 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 train_config["muon_grouped_gram_ns_fp32_byte_limit"] == 23 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 + + +def test_load_server_arguments_preserves_runner_compatibility_fields(tmp_path): + config_path = tmp_path / "server_config.yaml" + config_path.write_text( + yaml.safe_dump( + { + "model": { + "model_path": "Qwen/Qwen3-8B", + "record_routing_weights": False, + }, + "train": { + "enable_full_determinism": True, + "optimizer": "muon", + "cautious_weight_decay": True, + "muon_distributed_mode": "full_gradient", + "moe_grad_reduce_mode": "bf16_a2a_fp32_sum", + "output_dir": str(tmp_path / "outputs"), + }, + "lora": { + "lora_export_format": "sglang_shared_outer", + }, + } + ), + encoding="utf-8", + ) + + args = load_server_arguments(str(config_path)) + config = args.to_config_dict() + + assert config["model"]["record_routing_weights"] is False + assert config["train"]["enable_full_determinism"] is True + assert config["train"]["cautious_weight_decay"] is True + assert config["train"]["muon_distributed_mode"] == "full_gradient" + assert config["train"]["moe_grad_reduce_mode"] == "bf16_a2a_fp32_sum" + assert config["train"]["optimizer_kwargs"]["muon_distributed_mode"] == "full_gradient" + assert config["lora"]["lora_export_format"] == "sglang_shared_outer" diff --git a/tests/server/weight_sync/test_endpoint_routing.py b/tests/server/weight_sync/test_endpoint_routing.py index c16c71d9..9f337b4e 100644 --- a/tests/server/weight_sync/test_endpoint_routing.py +++ b/tests/server/weight_sync/test_endpoint_routing.py @@ -2,6 +2,8 @@ from unittest.mock import MagicMock, patch +import pytest +import requests import torch from xorl.server.weight_sync.backends.nccl_broadcast import EndpointInfo, NCCLWeightSynchronizer @@ -9,10 +11,13 @@ class FakeResponse: - def __init__(self, payload): + def __init__(self, payload, error: Exception | None = None): self._payload = payload + self._error = error def raise_for_status(self): + if self._error: + raise self._error return None def json(self): @@ -43,7 +48,34 @@ def test_endpoint_manager_health_check_uses_endpoint_port(self): 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) + session.get.assert_called_once_with("http://127.0.0.1:30000/model_info", timeout=60) + + def test_endpoint_manager_health_check_falls_back_to_v1_models(self): + session = MagicMock() + session.get.side_effect = [ + FakeResponse({}, requests.HTTPError("503 Server Error")), + FakeResponse({"data": [{"id": "Qwen/Qwen3-30B-A3B"}]}), + ] + 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() + + assert [call.args[0] for call in session.get.call_args_list] == [ + "http://127.0.0.1:30000/model_info", + "http://127.0.0.1:30000/v1/models", + ] + + def test_endpoint_manager_health_check_reports_all_failures(self): + session = MagicMock() + session.get.side_effect = requests.ConnectionError("connection refused") + manager = EndpointManager([{"host": "127.0.0.1", "port": 30000}]) + + with ( + patch("xorl.server.weight_sync.endpoint_manager._get_http_session", return_value=session), + pytest.raises(RuntimeError, match="/model_info.*/v1/models.*/health"), + ): + manager.health_check() def test_nccl_sync_init_uses_endpoint_port(self): session = MagicMock() diff --git a/tests/server/weight_sync/test_fp8_quantization.py b/tests/server/weight_sync/test_fp8_quantization.py new file mode 100644 index 00000000..7beca75c --- /dev/null +++ b/tests/server/weight_sync/test_fp8_quantization.py @@ -0,0 +1,483 @@ +import pytest +import torch + +from xorl.server.weight_sync.handler import WeightSyncHandler + + +def test_fp8_quantization_emits_cpu_weight_and_scale_tensors(): + name = "model.layers.0.mlp.gate_proj.weight" + tensor = torch.arange(32, dtype=torch.bfloat16).reshape(4, 8) + if torch.cuda.is_available(): + tensor = tensor.cuda() + + out = dict( + WeightSyncHandler._quantize_buffer_for_fp8( + [(name, tensor)], + quantization_config={ + "quant_method": "fp8", + "fmt": "e4m3", + "weight_block_size": [2, 4], + }, + target_device="cpu", + ) + ) + + assert set(out) == {name, "model.layers.0.mlp.gate_proj.weight_scale_inv"} + quantized = out[name] + scale = out["model.layers.0.mlp.gate_proj.weight_scale_inv"] + assert quantized.device.type == "cpu" + assert scale.device.type == "cpu" + assert quantized.dtype == torch.float8_e4m3fn + assert scale.dtype == torch.float32 + assert quantized.shape == (4, 8) + assert scale.shape == (2, 2) + assert torch.all(scale > 0) + + +def test_fp8_quantization_skips_default_non_projection_weights(): + tensor = torch.zeros(8, 4, dtype=torch.bfloat16) + out = WeightSyncHandler._quantize_buffer_for_fp8( + [("model.embed_tokens.weight", tensor)], + quantization_config={"quant_method": "fp8", "weight_block_size": [2, 4]}, + ) + + assert out == [("model.embed_tokens.weight", tensor)] + + +def test_fp8_quantization_includes_fused_mla_weight_by_default(): + name = "model.layers.0.self_attn.fused_qkv_a_proj_with_mqa.weight" + tensor = torch.zeros(8, 4, dtype=torch.bfloat16) + out = dict( + WeightSyncHandler._quantize_buffer_for_fp8( + [(name, tensor)], + quantization_config={"quant_method": "fp8", "weight_block_size": [2, 4]}, + ) + ) + + assert set(out) == {name, "model.layers.0.self_attn.fused_qkv_a_proj_with_mqa.weight_scale_inv"} + assert out[name].dtype == torch.float8_e4m3fn + + +@pytest.mark.parametrize("projection", ["in_proj_qkv", "in_proj_z", "in_proj_b", "in_proj_a"]) +def test_fp8_quantization_includes_qwen_linear_attention_packed_weights_by_default(projection): + name = f"model.layers.0.linear_attn.{projection}.weight" + tensor = torch.zeros(8, 4, dtype=torch.bfloat16) + out = dict( + WeightSyncHandler._quantize_buffer_for_fp8( + [(name, tensor)], + quantization_config={"quant_method": "fp8", "weight_block_size": [2, 4]}, + ) + ) + + assert set(out) == {name, f"model.layers.0.linear_attn.{projection}.weight_scale_inv"} + assert out[name].dtype == torch.float8_e4m3fn + + +def test_fp8_quantization_respects_modules_to_not_convert(): + name = "model.layers.0.mlp.gate_proj.weight" + tensor = torch.zeros(8, 4, dtype=torch.bfloat16) + out = WeightSyncHandler._quantize_buffer_for_fp8( + [(name, tensor)], + quantization_config={ + "quant_method": "fp8", + "weight_block_size": [2, 4], + "modules_to_not_convert": ["model.layers.0.mlp.gate_proj"], + }, + ) + + assert out == [(name, tensor)] + + +def test_fp8_quantization_can_detect_contiguous_expert_slice_groups(): + stack = torch.zeros(3, 4, 8, dtype=torch.bfloat16) + + assert WeightSyncHandler._can_group_fp8_tensor(stack[0], stack[1], 1) + assert WeightSyncHandler._can_group_fp8_tensor(stack[0], stack[2], 2) + assert not WeightSyncHandler._can_group_fp8_tensor(stack[0], stack[2], 1) + + +def test_fp8_stack_quantization_matches_single_tensor_quantization(): + stack = torch.arange(3 * 4 * 8, dtype=torch.bfloat16).reshape(3, 4, 8) + kwargs = { + "fp8_dtype": torch.float8_e4m3fn, + "fp8_max": torch.finfo(torch.float8_e4m3fn).max, + "block_size_row": 2, + "block_size_col": 4, + "target_device": "cpu", + "phase_s": {}, + "phase_prefix": "test_fp8", + } + + quantized_stack, scale_stack = WeightSyncHandler._quantize_fp8_stack(stack, **kwargs) + + for idx in range(stack.shape[0]): + quantized, scale = WeightSyncHandler._quantize_single_fp8_tensor(stack[idx], **kwargs) + assert torch.equal(quantized_stack[idx].float(), quantized.float()) + assert torch.equal(scale_stack[idx], scale) + + +def test_fp8_quantization_skips_already_quantized_weights(): + name = "model.layers.0.mlp.gate_proj.weight" + tensor = torch.zeros(4, 8, dtype=torch.float8_e4m3fn) + + out = WeightSyncHandler._quantize_buffer_for_fp8( + [(name, tensor)], + quantization_config={"quant_method": "fp8", "weight_block_size": [2, 4]}, + ) + + assert out == [(name, tensor)] + + +def test_fp8_cpu_expert_projection_quantization_emits_hf_weights_and_scales(): + local_data = torch.arange(2 * 4 * 8, dtype=torch.bfloat16).reshape(2, 4, 8) + phase_s = {} + + out, original_bytes = WeightSyncHandler._quantize_ep_expert_projection_for_fp8_cpu( + local_data, + full_prefix="model.layers.0.mlp.experts", + proj_name="gate_proj", + ep_rank=1, + quantization_config={"quant_method": "fp8", "fmt": "e4m3", "weight_block_size": [2, 4]}, + phase_s=phase_s, + ) + out_by_name = dict(out) + + assert original_bytes == local_data.numel() * local_data.element_size() + assert set(out_by_name) == { + "model.layers.0.mlp.experts.2.gate_proj.weight", + "model.layers.0.mlp.experts.2.gate_proj.weight_scale_inv", + "model.layers.0.mlp.experts.3.gate_proj.weight", + "model.layers.0.mlp.experts.3.gate_proj.weight_scale_inv", + } + assert out_by_name["model.layers.0.mlp.experts.2.gate_proj.weight"].shape == (8, 4) + assert out_by_name["model.layers.0.mlp.experts.2.gate_proj.weight"].dtype == torch.float8_e4m3fn + assert out_by_name["model.layers.0.mlp.experts.2.gate_proj.weight_scale_inv"].shape == (4, 1) + assert phase_s["direct_ep_fp8_cpu_transpose_s"] >= 0 + + +def test_fp8_cpu_expert_projection_can_defer_quantization(): + local_data = torch.arange(2 * 4 * 8, dtype=torch.bfloat16).reshape(2, 4, 8) + phase_s = {} + + out, original_bytes = WeightSyncHandler._format_ep_expert_projection_for_fp8_cpu( + local_data, + full_prefix="model.layers.0.mlp.experts", + proj_name="gate_proj", + ep_rank=1, + phase_s=phase_s, + ) + out_by_name = dict(out) + + assert original_bytes == local_data.numel() * local_data.element_size() + assert set(out_by_name) == { + "model.layers.0.mlp.experts.2.gate_proj.weight", + "model.layers.0.mlp.experts.3.gate_proj.weight", + } + assert out_by_name["model.layers.0.mlp.experts.2.gate_proj.weight"].shape == (8, 4) + assert out_by_name["model.layers.0.mlp.experts.2.gate_proj.weight"].dtype == torch.bfloat16 + assert out_by_name["model.layers.0.mlp.experts.2.gate_proj.weight"].device.type == "cpu" + assert phase_s["direct_ep_fp8_source_copy_s"] >= 0 + assert phase_s["direct_ep_fp8_cpu_transpose_s"] >= 0 + + +def test_fp8_cpu_workspace_stages_quantizes_and_reuses_storage(monkeypatch): + monkeypatch.setenv("XORL_P2P_FP8_CPU_WORKSPACE", "1") + monkeypatch.setenv("XORL_P2P_FP8_CPU_WORKSPACE_PINNED", "0") + monkeypatch.setenv("XORL_P2P_FP8_CPU_WORKSPACE_MIN_CAPACITY", "2") + handler = WeightSyncHandler(rank=0, world_size=1, trainer=None) + local_data = torch.arange(2 * 4 * 8, dtype=torch.bfloat16).reshape(2, 4, 8) + phase_s = {} + quantization_config = {"quant_method": "fp8", "fmt": "e4m3", "weight_block_size": [2, 4]} + + records, original_bytes = handler._stage_ep_expert_projection_for_fp8_cpu_workspace( + local_data, + full_prefix="model.layers.0.mlp.experts", + proj_name="gate_proj", + ep_rank=1, + quantization_config=quantization_config, + phase_s=phase_s, + ) + + assert original_bytes == local_data.numel() * local_data.element_size() + assert [name for name, _, _ in records] == [ + "model.layers.0.mlp.experts.2.gate_proj.weight", + "model.layers.0.mlp.experts.3.gate_proj.weight", + ] + workspace = handler._fp8_cpu_workspaces[records[0][1]] + assert torch.equal(workspace["input"][:2], local_data.permute(0, 2, 1).contiguous()) + input_ptr = workspace["input"].data_ptr() + + out = handler._quantize_fp8_cpu_workspace_records( + records, + quantization_config=quantization_config, + phase_s=phase_s, + phase_prefix="test_fp8", + ) + assert [name for name, _ in out] == [ + "model.layers.0.mlp.experts.2.gate_proj.weight", + "model.layers.0.mlp.experts.2.gate_proj.weight_scale_inv", + "model.layers.0.mlp.experts.3.gate_proj.weight", + "model.layers.0.mlp.experts.3.gate_proj.weight_scale_inv", + ] + out_by_name = dict(out) + assert out_by_name["model.layers.0.mlp.experts.2.gate_proj.weight"].shape == (8, 4) + assert out_by_name["model.layers.0.mlp.experts.2.gate_proj.weight"].dtype == torch.float8_e4m3fn + assert out_by_name["model.layers.0.mlp.experts.2.gate_proj.weight_scale_inv"].shape == (4, 1) + assert phase_s["direct_ep_fp8_workspace_alloc_s"] >= 0 + assert phase_s["direct_ep_fp8_workspace_copy_s"] >= 0 + assert phase_s["test_fp8_float_s"] >= 0 + assert phase_s["test_fp8_reduce_s"] >= 0 + assert phase_s["test_fp8_cast_s"] >= 0 + + handler._reset_fp8_cpu_workspace_usage() + records, _ = handler._stage_ep_expert_projection_for_fp8_cpu_workspace( + local_data, + full_prefix="model.layers.0.mlp.experts", + proj_name="gate_proj", + ep_rank=1, + quantization_config=quantization_config, + phase_s=phase_s, + ) + assert handler._fp8_cpu_workspaces[records[0][1]]["input"].data_ptr() == input_ptr + + +def test_fp8_cpu_workspace_streams_quantized_chunks(monkeypatch): + class RecordingBackend: + def __init__(self): + self.calls = [] + + def transfer_bucket(self, bucket, *, src_rank=0, flush_cache=False, weight_version=None): + self.calls.append( + { + "names": [name for name, _ in bucket], + "dtypes": [tensor.dtype for _, tensor in bucket], + "src_rank": src_rank, + "flush_cache": flush_cache, + "weight_version": weight_version, + } + ) + + monkeypatch.setenv("XORL_P2P_FP8_CPU_WORKSPACE", "1") + monkeypatch.setenv("XORL_P2P_FP8_CPU_WORKSPACE_PINNED", "0") + monkeypatch.setenv("XORL_P2P_FP8_CPU_WORKSPACE_MIN_CAPACITY", "4") + monkeypatch.setenv("XORL_P2P_FP8_CPU_WORKSPACE_STREAMING", "1") + monkeypatch.setenv("XORL_P2P_FP8_CPU_WORKSPACE_STREAM_BYTES", "96") + handler = WeightSyncHandler(rank=3, world_size=4, trainer=None) + backend = RecordingBackend() + local_data = torch.arange(4 * 4 * 8, dtype=torch.bfloat16).reshape(4, 4, 8) + phase_s = {} + quantization_config = {"quant_method": "fp8", "fmt": "e4m3", "weight_block_size": [2, 4]} + + records, _ = handler._stage_ep_expert_projection_for_fp8_cpu_workspace( + local_data, + full_prefix="model.layers.0.mlp.experts", + proj_name="gate_proj", + ep_rank=0, + quantization_config=quantization_config, + phase_s=phase_s, + ) + + num_buckets = handler._quantize_and_transfer_fp8_cpu_workspace_records( + backend, + records, + quantization_config=quantization_config, + bucket_size_bytes=96, + flush_cache=True, + weight_version="sync-1", + phase_s=phase_s, + phase_prefix="test_fp8", + ) + + assert num_buckets == 2 + assert len(backend.calls) == 2 + assert backend.calls[0]["src_rank"] == 3 + assert backend.calls[0]["flush_cache"] is False + assert backend.calls[0]["weight_version"] is None + assert backend.calls[1]["flush_cache"] is True + assert backend.calls[1]["weight_version"] == "sync-1" + assert backend.calls[0]["names"] == [ + "model.layers.0.mlp.experts.0.gate_proj.weight", + "model.layers.0.mlp.experts.0.gate_proj.weight_scale_inv", + "model.layers.0.mlp.experts.1.gate_proj.weight", + "model.layers.0.mlp.experts.1.gate_proj.weight_scale_inv", + ] + assert backend.calls[1]["dtypes"] == [ + torch.float8_e4m3fn, + torch.float32, + torch.float8_e4m3fn, + torch.float32, + ] + assert phase_s["test_fp8_float_s"] >= 0 + assert phase_s["test_fp8_reduce_s"] >= 0 + assert phase_s["test_fp8_cast_s"] >= 0 + assert phase_s["direct_ep_backend_s"] >= 0 + assert phase_s["direct_ep_fp8_workspace_stream_wait_s"] >= 0 + + +def test_fp8_cpu_workspace_flush_resets_used_capacity(monkeypatch): + class RecordingBackend: + def __init__(self): + self.calls = [] + + def transfer_bucket(self, bucket, *, src_rank=0, flush_cache=False, weight_version=None): + self.calls.append( + { + "names": [name for name, _ in bucket], + "flush_cache": flush_cache, + "weight_version": weight_version, + } + ) + + monkeypatch.setenv("XORL_P2P_FP8_CPU_WORKSPACE", "1") + monkeypatch.setenv("XORL_P2P_FP8_CPU_WORKSPACE_PINNED", "0") + monkeypatch.setenv("XORL_P2P_FP8_CPU_WORKSPACE_MIN_CAPACITY", "2") + handler = WeightSyncHandler(rank=0, world_size=1, trainer=None) + backend = RecordingBackend() + quantization_config = {"quant_method": "fp8", "fmt": "e4m3", "weight_block_size": [2, 4]} + local_data = torch.arange(2 * 4 * 8, dtype=torch.bfloat16).reshape(2, 4, 8) + phase_s = {} + + records, original_bytes = handler._stage_ep_expert_projection_for_fp8_cpu_workspace( + local_data, + full_prefix="model.layers.0.mlp.experts", + proj_name="gate_proj", + ep_rank=0, + quantization_config=quantization_config, + phase_s=phase_s, + ) + handler._pending_moe_cpu_workspace_records.extend(records) + handler._pending_moe_bucket_bytes += original_bytes + workspace = handler._fp8_cpu_workspaces[records[0][1]] + input_ptr = workspace["input"].data_ptr() + + _, _, num_buckets = handler._flush_pending_moe_bucket( + backend, + flush_cache=False, + weight_version=None, + quantization=quantization_config, + bucket_size_bytes=1024, + phase_s=phase_s, + ) + + assert num_buckets == 1 + assert backend.calls[0]["flush_cache"] is False + assert backend.calls[0]["weight_version"] is None + assert handler._pending_moe_cpu_workspace_records == [] + assert handler._pending_moe_bucket_bytes == 0 + assert workspace["used"] == 0 + + records, _ = handler._stage_ep_expert_projection_for_fp8_cpu_workspace( + local_data, + full_prefix="model.layers.1.mlp.experts", + proj_name="gate_proj", + ep_rank=0, + quantization_config=quantization_config, + phase_s=phase_s, + ) + assert handler._fp8_cpu_workspaces[records[0][1]]["input"].data_ptr() == input_ptr + assert [index for _, _, index in records] == [0, 1] + + +def test_empty_moe_final_flush_preserves_p2p_completion_metadata(): + class Config: + def __init__(self): + self.backend_config = {} + + class Backend: + def __init__(self): + self.config = Config() + + handler = WeightSyncHandler(rank=0, world_size=1, trainer=None) + backend = Backend() + + _, _, num_buckets = handler._flush_pending_moe_bucket( + backend, + flush_cache=True, + weight_version="sync-2", + quantization={"quant_method": "fp8"}, + bucket_size_bytes=1024, + phase_s={}, + ) + + assert num_buckets == 0 + assert backend.config.backend_config["flush_cache"] is True + assert backend.config.backend_config["weight_version"] == "sync-2" + + +def test_fp8_cpu_expert_projection_respects_modules_to_not_convert(): + local_data = torch.arange(2 * 4 * 8, dtype=torch.bfloat16).reshape(2, 4, 8) + + out, _ = WeightSyncHandler._quantize_ep_expert_projection_for_fp8_cpu( + local_data, + full_prefix="model.layers.0.mlp.experts", + proj_name="gate_proj", + ep_rank=0, + quantization_config={ + "quant_method": "fp8", + "fmt": "e4m3", + "weight_block_size": [2, 4], + "modules_to_not_convert": ["model.layers.0.mlp.experts"], + }, + phase_s={}, + ) + out_by_name = dict(out) + + assert set(out_by_name) == { + "model.layers.0.mlp.experts.0.gate_proj.weight", + "model.layers.0.mlp.experts.1.gate_proj.weight", + } + assert out_by_name["model.layers.0.mlp.experts.0.gate_proj.weight"].dtype == torch.bfloat16 + assert out_by_name["model.layers.0.mlp.experts.0.gate_proj.weight"].device.type == "cpu" + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_fp8_gpu_stack_quantization_returns_cpu_tensors(monkeypatch): + monkeypatch.setenv("XORL_P2P_FP8_QUANTIZE_DEVICE", "gpu") + stack = torch.arange(2 * 128 * 128, dtype=torch.bfloat16, device="cuda").reshape(2, 128, 128) + phase_s = {} + + quantized, scale = WeightSyncHandler._quantize_fp8_stack( + stack, + fp8_dtype=torch.float8_e4m3fn, + fp8_max=torch.finfo(torch.float8_e4m3fn).max, + block_size_row=128, + block_size_col=128, + target_device="cpu", + phase_s=phase_s, + phase_prefix="test_fp8", + ) + + assert quantized.device.type == "cpu" + assert scale.device.type == "cpu" + assert quantized.dtype == torch.float8_e4m3fn + assert scale.shape == (2, 1, 1) + assert phase_s["test_fp8_gpu_quant_s"] >= 0 + assert phase_s["test_fp8_gpu_output_copy_s"] >= 0 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_fp8_gpu_expert_projection_respects_modules_to_not_convert(monkeypatch): + monkeypatch.setenv("XORL_P2P_FP8_QUANTIZE_DEVICE", "gpu") + local_data = torch.arange(2 * 128 * 128, dtype=torch.bfloat16, device="cuda").reshape(2, 128, 128) + + out, _ = WeightSyncHandler._quantize_ep_expert_projection_for_fp8_gpu_to_cpu( + local_data, + full_prefix="model.layers.0.mlp.experts", + proj_name="gate_proj", + ep_rank=0, + quantization_config={ + "quant_method": "fp8", + "fmt": "e4m3", + "weight_block_size": [128, 128], + "modules_to_not_convert": ["model.layers.0.mlp.experts"], + }, + phase_s={}, + ) + out_by_name = dict(out) + + assert set(out_by_name) == { + "model.layers.0.mlp.experts.0.gate_proj.weight", + "model.layers.0.mlp.experts.1.gate_proj.weight", + } + assert out_by_name["model.layers.0.mlp.experts.0.gate_proj.weight"].dtype == torch.bfloat16 + assert out_by_name["model.layers.0.mlp.experts.0.gate_proj.weight"].device.type == "cpu" diff --git a/tests/server/weight_sync/test_handler_config.py b/tests/server/weight_sync/test_handler_config.py new file mode 100644 index 00000000..3dadeeb4 --- /dev/null +++ b/tests/server/weight_sync/test_handler_config.py @@ -0,0 +1,340 @@ +from types import SimpleNamespace + +import torch + +from xorl.server.weight_sync.handler import ( + _DEFAULT_MOE_BUCKET_BYTES, + _DEFAULT_P2P_MOE_BUCKET_BYTES, + WeightSyncHandler, + _moe_bucket_size_bytes, + _p2p_direct_ep_sender_ep_ranks, + _p2p_direct_ep_sender_ranks, + _should_collect_ep_moe_tensors, +) + + +def test_moe_bucket_default_is_backend_specific(monkeypatch): + monkeypatch.delenv("XORL_WEIGHT_SYNC_MOE_BUCKET_BYTES", raising=False) + monkeypatch.delenv("XORL_WEIGHT_SYNC_BUCKET_BYTES", raising=False) + + assert _moe_bucket_size_bytes("nccl_broadcast") == _DEFAULT_MOE_BUCKET_BYTES + assert _moe_bucket_size_bytes("p2p") == _DEFAULT_P2P_MOE_BUCKET_BYTES + + +def test_legacy_moe_bucket_env_override_is_explicit(monkeypatch): + monkeypatch.delenv("XORL_WEIGHT_SYNC_MOE_BUCKET_BYTES", raising=False) + monkeypatch.setenv("XORL_WEIGHT_SYNC_BUCKET_BYTES", str(123 * 1024 * 1024)) + + assert _moe_bucket_size_bytes("nccl_broadcast") == 123 * 1024 * 1024 + assert _moe_bucket_size_bytes("p2p") == 123 * 1024 * 1024 + + +def test_dedicated_moe_bucket_env_override_takes_precedence(monkeypatch): + monkeypatch.setenv("XORL_WEIGHT_SYNC_BUCKET_BYTES", str(123 * 1024 * 1024)) + monkeypatch.setenv("XORL_WEIGHT_SYNC_MOE_BUCKET_BYTES", str(456 * 1024 * 1024)) + + assert _moe_bucket_size_bytes("nccl_broadcast") == 456 * 1024 * 1024 + assert _moe_bucket_size_bytes("p2p") == 456 * 1024 * 1024 + + +def test_p2p_fp8_sync_requires_receiver_post_process(): + assert WeightSyncHandler._p2p_requires_post_process_weights({"quant_method": "fp8"}) is True + assert WeightSyncHandler._p2p_requires_post_process_weights({"quant_method": "awq"}) is False + assert WeightSyncHandler._p2p_requires_post_process_weights(None) is False + + +def test_chunk_buffer_by_bytes_splits_large_dense_items(): + items = [ + ("a", torch.zeros(4, dtype=torch.float32)), + ("b", torch.zeros(8, dtype=torch.float32)), + ("c", torch.zeros(4, dtype=torch.float32)), + ] + + chunks = WeightSyncHandler._chunk_buffer_by_bytes(items, bucket_size_bytes=32) + + assert [[name for name, _ in chunk] for chunk in chunks] == [["a"], ["b"], ["c"]] + + +def test_would_exceed_bucket_cap_flushes_before_oversized_append(): + assert WeightSyncHandler._would_exceed_bucket_cap(900, 200, 1024) is True + assert WeightSyncHandler._would_exceed_bucket_cap(0, 2048, 1024) is False + assert WeightSyncHandler._would_exceed_bucket_cap(800, 224, 1024) is False + + +def test_p2p_direct_ep_sender_ranks_default_to_first_ep_fsdp_replica(monkeypatch): + monkeypatch.delenv("XORL_P2P_DIRECT_EP_REPLICA_STRATEGY", raising=False) + ps = SimpleNamespace( + ep_enabled=True, + ep_fsdp_device_mesh=SimpleNamespace( + mesh=torch.tensor( + [ + [0, 8, 16, 24], + [1, 9, 17, 25], + [2, 10, 18, 26], + [3, 11, 19, 27], + [4, 12, 20, 28], + [5, 13, 21, 29], + [6, 14, 22, 30], + [7, 15, 23, 31], + ] + ) + ), + ) + + assert _p2p_direct_ep_sender_ranks(ps, 32) == tuple(range(8)) + + +def test_p2p_direct_ep_sender_ranks_can_round_robin_replicas(monkeypatch): + monkeypatch.setenv("XORL_P2P_DIRECT_EP_REPLICA_STRATEGY", "round_robin") + ps = SimpleNamespace( + ep_enabled=True, + ep_fsdp_device_mesh=SimpleNamespace( + mesh=torch.tensor( + [ + [0, 8, 16, 24], + [1, 9, 17, 25], + [2, 10, 18, 26], + [3, 11, 19, 27], + [4, 12, 20, 28], + [5, 13, 21, 29], + [6, 14, 22, 30], + [7, 15, 23, 31], + ] + ) + ), + ) + + assert _p2p_direct_ep_sender_ranks(ps, 32) == (0, 4, 9, 13, 18, 22, 27, 31) + + +def test_p2p_direct_ep_sender_ep_ranks_tracks_selected_replicas(monkeypatch): + monkeypatch.setenv("XORL_P2P_DIRECT_EP_REPLICA_STRATEGY", "round_robin") + ps = SimpleNamespace( + ep_enabled=True, + ep_size=8, + ep_fsdp_device_mesh=SimpleNamespace( + mesh=torch.tensor( + [ + [0, 8, 16, 24], + [1, 9, 17, 25], + [2, 10, 18, 26], + [3, 11, 19, 27], + [4, 12, 20, 28], + [5, 13, 21, 29], + [6, 14, 22, 30], + [7, 15, 23, 31], + ] + ) + ), + ) + sender_ranks = _p2p_direct_ep_sender_ranks(ps, 32) + + assert _p2p_direct_ep_sender_ep_ranks(ps, sender_ranks, 32) == ( + (0, 0), + (4, 4), + (9, 1), + (13, 5), + (18, 2), + (22, 6), + (27, 3), + (31, 7), + ) + + +def test_p2p_explicit_direct_ep_non_senders_skip_moe_tensor_collection(): + backend = SimpleNamespace( + supports_direct_ep_transfer=True, + has_explicit_sender_ranks=True, + ) + + assert _should_collect_ep_moe_tensors("p2p", backend, is_sender=False) is False + assert _should_collect_ep_moe_tensors("p2p", backend, is_sender=True) is True + + +def test_moe_tensor_collection_is_kept_for_non_explicit_paths(): + direct_backend = SimpleNamespace( + supports_direct_ep_transfer=True, + has_explicit_sender_ranks=False, + ) + nccl_backend = SimpleNamespace( + supports_direct_ep_transfer=False, + has_explicit_sender_ranks=False, + ) + + assert _should_collect_ep_moe_tensors("p2p", direct_backend, is_sender=False) is True + assert _should_collect_ep_moe_tensors("nccl_broadcast", nccl_backend, is_sender=False) is True + + +def test_extract_params_can_defer_tied_weight_duplicate_for_p2p(): + class Root(torch.nn.Module): + _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} + + def __init__(self): + super().__init__() + self.model = torch.nn.Module() + self.model.embed_tokens = torch.nn.Embedding(4, 3) + self.lm_head = torch.nn.Linear(3, 4, bias=False) + self.lm_head.weight = self.model.embed_tokens.weight + + aliases = {} + + class FakeDTensor: + pass + + buffer = WeightSyncHandler._extract_params_for_sync( + Root(), + "(root)", + FakeDTensor, + emit_tied_weight_duplicates=False, + tied_weight_aliases=aliases, + ) + + assert [name for name, _ in buffer] == ["model.embed_tokens.weight"] + assert aliases == {"lm_head.weight": "model.embed_tokens.weight"} + + +def test_extract_params_include_param_skips_unowned_dense_params(): + class Root(torch.nn.Module): + def __init__(self): + super().__init__() + self.keep = torch.nn.Linear(3, 4, bias=False) + self.skip = torch.nn.Linear(3, 4, bias=False) + + class FakeDTensor: + pass + + buffer = WeightSyncHandler._extract_params_for_sync( + Root(), + "(root)", + FakeDTensor, + include_param=lambda name: name == "keep.weight", + ) + + assert [name for name, _ in buffer] == ["keep.weight"] + + +def test_extract_params_skips_tied_weight_seen_in_prior_module_for_p2p(): + class Root(torch.nn.Module): + _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} + + def __init__(self): + super().__init__() + self.model = torch.nn.Module() + self.model.embed_tokens = torch.nn.Embedding(4, 3) + self.lm_head = torch.nn.Linear(3, 4, bias=False) + self.lm_head.weight = self.model.embed_tokens.weight + + class FakeDTensor: + pass + + root = Root() + aliases = {} + + root_buffer = WeightSyncHandler._extract_params_for_sync( + root, + "(root)", + FakeDTensor, + emit_tied_weight_duplicates=False, + tied_weight_aliases=aliases, + ) + lm_head_buffer = WeightSyncHandler._extract_params_for_sync( + root.lm_head, + "lm_head", + FakeDTensor, + emit_tied_weight_duplicates=False, + tied_weight_aliases=aliases, + ) + + assert [name for name, _ in root_buffer] == ["model.embed_tokens.weight"] + assert lm_head_buffer == [] + assert aliases == {"lm_head.weight": "model.embed_tokens.weight"} + + +def test_extract_params_does_not_defer_declared_tie_when_parameters_differ(): + class Root(torch.nn.Module): + _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} + + def __init__(self): + super().__init__() + self.model = torch.nn.Module() + self.model.embed_tokens = torch.nn.Embedding(4, 3) + self.lm_head = torch.nn.Linear(3, 4, bias=False) + + def named_parameters(self, *args, **kwargs): + yield "model.embed_tokens.weight", self.model.embed_tokens.weight + + class FakeDTensor: + pass + + aliases = {} + buffer = WeightSyncHandler._extract_params_for_sync( + Root(), + "(root)", + FakeDTensor, + emit_tied_weight_duplicates=False, + tied_weight_aliases=aliases, + ) + + assert [name for name, _ in buffer] == ["model.embed_tokens.weight"] + assert aliases == {} + + +def test_extract_params_does_not_infer_tied_weight_from_storage_only(): + class FakeDTensor: + pass + + source = torch.nn.Linear(3, 4, bias=False) + alias = torch.nn.Linear(3, 4, bias=False) + alias.weight = source.weight + + aliases = {} + source_buffer = WeightSyncHandler._extract_params_for_sync( + source, + "source", + FakeDTensor, + emit_tied_weight_duplicates=False, + tied_weight_aliases=aliases, + ) + alias_buffer = WeightSyncHandler._extract_params_for_sync( + alias, + "alias", + FakeDTensor, + emit_tied_weight_duplicates=False, + tied_weight_aliases=aliases, + ) + + assert [name for name, _ in source_buffer] == ["source.weight"] + assert [name for name, _ in alias_buffer] == ["alias.weight"] + assert aliases == {} + + +def test_unfuse_for_inference_fuses_deepseek_kimi_mla_a_projection_for_sglang(): + config = SimpleNamespace( + hidden_size=8, + num_attention_heads=2, + q_lora_rank=3, + layer_types=[], + ) + model = SimpleNamespace(config=config) + q_a = torch.arange(3 * 8, dtype=torch.bfloat16).reshape(3, 8) + kv_a = torch.arange(5 * 8, dtype=torch.bfloat16).reshape(5, 8) + q_b = torch.ones(4, 3, dtype=torch.bfloat16) + + remapped = dict( + WeightSyncHandler._unfuse_for_inference( + [ + ("model.layers.0.self_attn.q_a_proj.weight", q_a), + ("model.layers.0.self_attn.kv_a_proj_with_mqa.weight", kv_a), + ("model.layers.0.self_attn.q_b_proj.weight", q_b), + ], + model, + ) + ) + + assert "model.layers.0.self_attn.q_a_proj.weight" not in remapped + assert "model.layers.0.self_attn.kv_a_proj_with_mqa.weight" not in remapped + torch.testing.assert_close( + remapped["model.layers.0.self_attn.fused_qkv_a_proj_with_mqa.weight"], + torch.cat([q_a, kv_a], dim=0), + ) + torch.testing.assert_close(remapped["model.layers.0.self_attn.q_b_proj.weight"], q_b) diff --git a/tests/server/weight_sync/test_moe_runtime_scaling.py b/tests/server/weight_sync/test_moe_runtime_scaling.py new file mode 100644 index 00000000..af45daf6 --- /dev/null +++ b/tests/server/weight_sync/test_moe_runtime_scaling.py @@ -0,0 +1,78 @@ +"""Regression tests for runtime-rank MoE LoRA weight sync.""" + +import pytest +import torch + +from xorl.server.weight_sync.handler import WeightSyncHandler + + +pytestmark = [pytest.mark.cpu, pytest.mark.server] + + +class _FakeRuntimeRankMoeModule: + def __init__(self) -> None: + self.num_local_experts = 1 + self.hidden_size = 1 + self.intermediate_size = 1 + self.active_r = 2 + self.scaling = 4.0 + + def _active_scaling(self) -> float: + return 2.0 + + def dequantize_expert(self, _proj_name: str, _expert_idx: int, K: int, N: int) -> torch.Tensor: + return torch.zeros((K, N), dtype=torch.float32) + + +def _runtime_rank_lora_params() -> dict[str, torch.Tensor]: + lora_A = torch.tensor([[[1.0, 1.0, 0.0, 0.0]]], dtype=torch.float32) + lora_B = torch.tensor([[[1.0], [1.0], [0.0], [0.0]]], dtype=torch.float32) + return { + "gate_proj_lora_A": lora_A.clone(), + "gate_proj_lora_B": lora_B.clone(), + "up_proj_lora_A": lora_A.clone(), + "up_proj_lora_B": lora_B.clone(), + "down_proj_lora_A": lora_A.clone(), + "down_proj_lora_B": lora_B.clone(), + } + + +def test_compute_moe_lora_delta_uses_active_runtime_scaling(): + mod = _FakeRuntimeRankMoeModule() + params = _runtime_rank_lora_params() + + delta = WeightSyncHandler._compute_moe_lora_delta( + mod, + params["gate_proj_lora_A"], + params["gate_proj_lora_B"], + expert_idx=0, + ) + + assert delta.shape == (1, 1) + assert delta.item() == pytest.approx(4.0) + + +def test_compute_moe_experts_buffer_uses_runtime_scaled_delta(): + mod = _FakeRuntimeRankMoeModule() + handler = object.__new__(WeightSyncHandler) + handler.rank = 0 + + ctx = { + "module": mod, + "prefix": "model.layers.0.mlp.experts", + "lora_params": _runtime_rank_lora_params(), + } + + items = WeightSyncHandler._compute_moe_experts_buffer(handler, ctx) + + assert len(items) == 3 + assert [name for name, _ in items] == [ + "model.layers.0.mlp.experts.0.gate_proj.weight", + "model.layers.0.mlp.experts.0.up_proj.weight", + "model.layers.0.mlp.experts.0.down_proj.weight", + ] + for _, tensor in items: + assert tensor.dtype == torch.bfloat16 + assert tensor.shape == (1, 1) + assert tensor.float().item() == pytest.approx(4.0) + assert ctx["lora_params"] is None diff --git a/tests/server/weight_sync/test_nccl_broadcast_lifecycle.py b/tests/server/weight_sync/test_nccl_broadcast_lifecycle.py index 93c8d727..d1915a54 100644 --- a/tests/server/weight_sync/test_nccl_broadcast_lifecycle.py +++ b/tests/server/weight_sync/test_nccl_broadcast_lifecycle.py @@ -33,7 +33,7 @@ class _StickyTCPStore: open_ports = set() next_ephemeral_port = 31000 - def __init__(self, host_name, port, world_size, is_master, timeout): + def __init__(self, host_name, port, world_size, is_master, timeout, **_kwargs): del host_name, world_size, is_master, timeout if port == 0: port = self.next_ephemeral_port diff --git a/tests/server/weight_sync/test_p2p_async_api.py b/tests/server/weight_sync/test_p2p_async_api.py new file mode 100644 index 00000000..9e411f94 --- /dev/null +++ b/tests/server/weight_sync/test_p2p_async_api.py @@ -0,0 +1,201 @@ +import pytest + +from xorl.server.weight_sync.backends import p2p +from xorl.server.weight_sync.backends.base import EndpointConfig, TransportConfig +from xorl.server.weight_sync.backends.p2p import P2PTransportBackend, _BucketTiming, _do_async_transfer + + +class DoneEvent: + def synchronize(self): + return None + + +class FakeAsyncEngine: + def __init__(self, statuses): + self.statuses = list(statuses) + self.submitted = [] + + def batch_transfer_async_write(self, session_id, src_ptrs, peer_ptrs, lengths): + self.submitted.append((session_id, src_ptrs, peer_ptrs, lengths)) + return len(self.submitted) + + def get_batch_transfer_status(self, _bids): + if self.statuses: + return self.statuses.pop(0) + return 1 + + +class FakeEngineWrapper: + def __init__(self, statuses): + self.engine = FakeAsyncEngine(statuses) + self.sync_submitted = [] + + def batch_transfer_sync(self, session_id, src_ptrs, peer_ptrs, lengths): + self.sync_submitted.append((session_id, src_ptrs, peer_ptrs, lengths)) + return 0 + + +class FakeLocalEngine: + def get_session_id(self): + return "sender-session" + + def get_ib_device(self): + return "mlx5_0" + + +class FakePrepareResponse: + status_code = 200 + text = "" + + def json(self): + return { + "success": True, + "tensor_map": { + "model.embed_tokens.weight": [ + { + "hf_name": "model.embed_tokens.weight", + "tp_rank": 0, + "slice": [[0, 1], [0, 1]], + "full_shape": [1, 1], + "ptr": 1234, + "nbytes": 2, + "session_id": "receiver-session", + } + ] + }, + "receiver_transfer_engine_infos": [{"session_id": "receiver-session"}], + } + + +def _session_entries(src_ptrs, peer_ptrs, lengths): + return (src_ptrs, peer_ptrs, lengths, []) + + +def test_async_api_success_status_zero_completes(monkeypatch): + monkeypatch.setenv("XORL_P2P_USE_ASYNC_API", "1") + wrapper = FakeEngineWrapper([0]) + timing = _BucketTiming() + + _do_async_transfer( + engine_wrapper=wrapper, + copy_done_event=DoneEvent(), + by_session={"session-a": _session_entries([1], [2], [128 * 1024 * 1024])}, + small_session_data={}, + session_debug_info={"session-a": {"world_rank": 0}}, + small_register_ptrs=[], + small_register_lens=[], + chunk=1, + use_async_api=True, + timing=timing, + bucket_idx=1, + slice_holds=[], + src_view_holds=[], + log_bucket_details=False, + ) + + assert wrapper.engine.submitted == [("session-a", [1], [2], [128 * 1024 * 1024])] + assert wrapper.sync_submitted == [] + assert timing.transfer_s >= 0 + + +def test_async_api_uses_sync_fallback_for_medium_chunks(monkeypatch): + monkeypatch.setenv("XORL_P2P_USE_ASYNC_API", "1") + wrapper = FakeEngineWrapper([0]) + timing = _BucketTiming() + + _do_async_transfer( + engine_wrapper=wrapper, + copy_done_event=DoneEvent(), + by_session={"session-a": _session_entries([1], [2], [12 * 1024 * 1024])}, + small_session_data={}, + session_debug_info={"session-a": {"world_rank": 0}}, + small_register_ptrs=[], + small_register_lens=[], + chunk=1, + use_async_api=True, + timing=timing, + bucket_idx=1, + slice_holds=[], + src_view_holds=[], + log_bucket_details=False, + ) + + assert wrapper.engine.submitted == [] + assert wrapper.sync_submitted == [("session-a", [1], [2], [12 * 1024 * 1024])] + assert timing.transfer_s >= 0 + + +def test_async_api_min_bytes_env_controls_cutoff(monkeypatch): + monkeypatch.setenv("XORL_P2P_USE_ASYNC_API", "1") + monkeypatch.setenv("XORL_P2P_ASYNC_MIN_BYTES", str(8 * 1024 * 1024)) + wrapper = FakeEngineWrapper([0]) + timing = _BucketTiming() + + _do_async_transfer( + engine_wrapper=wrapper, + copy_done_event=DoneEvent(), + by_session={"session-a": _session_entries([1], [2], [12 * 1024 * 1024])}, + small_session_data={}, + session_debug_info={"session-a": {"world_rank": 0}}, + small_register_ptrs=[], + small_register_lens=[], + chunk=1, + use_async_api=True, + timing=timing, + bucket_idx=1, + slice_holds=[], + src_view_holds=[], + log_bucket_details=False, + ) + + assert wrapper.engine.submitted == [("session-a", [1], [2], [12 * 1024 * 1024])] + assert wrapper.sync_submitted == [] + assert timing.transfer_s >= 0 + + +def test_async_api_status_poll_timeout(monkeypatch): + monkeypatch.setenv("XORL_P2P_USE_ASYNC_API", "1") + monkeypatch.setenv("XORL_P2P_ASYNC_STATUS_TIMEOUT_S", "0.001") + wrapper = FakeEngineWrapper([1]) + + with pytest.raises(RuntimeError, match="async transfer status poll timed out"): + _do_async_transfer( + engine_wrapper=wrapper, + copy_done_event=DoneEvent(), + by_session={"session-a": _session_entries([1], [2], [128 * 1024 * 1024])}, + small_session_data={}, + session_debug_info={"session-a": {"world_rank": 0}}, + small_register_ptrs=[], + small_register_lens=[], + chunk=1, + use_async_api=True, + timing=_BucketTiming(), + bucket_idx=7, + slice_holds=[], + src_view_holds=[], + log_bucket_details=False, + ) + + +def test_prepare_uses_prepare_timeout_env(monkeypatch): + monkeypatch.setenv("XORL_P2P_PREPARE_TIMEOUT_S", "12.5") + seen = {} + + def fake_post(_url, *, json, timeout): + seen["payload"] = json + seen["timeout"] = timeout + return FakePrepareResponse() + + monkeypatch.setattr(p2p.requests, "post", fake_post) + backend = P2PTransportBackend( + TransportConfig( + endpoints=[EndpointConfig(host="receiver", port=30000, world_size=8)], + group_name="test-group", + training_rank=0, + ) + ) + monkeypatch.setattr(backend, "_make_local_engine", lambda: FakeLocalEngine()) + + assert backend._initialize_single_sender() + assert seen["timeout"] == 12.5 + assert seen["payload"]["transport"] == "p2p" diff --git a/tests/server/weight_sync/test_p2p_backend_protocol.py b/tests/server/weight_sync/test_p2p_backend_protocol.py new file mode 100644 index 00000000..40ff2798 --- /dev/null +++ b/tests/server/weight_sync/test_p2p_backend_protocol.py @@ -0,0 +1,2101 @@ +"""Protocol-level tests for the P2P weight transport backend. + +These tests exercise the xorl-side wire protocol against mocked SGLang HTTP +endpoints — no GPU, no Mooncake, no real network. They cover: + +* The shape of the ``/prepare_weights_update`` POST when ``transport=p2p``. +* Aggregating ``tensor_map`` and ``receiver_transfer_engine_infos`` across + multiple endpoints. +* Slicing the trainer's full HF tensor according to a locator's + ``slice``/``full_shape`` fields and computing the right Mooncake address. +* The shape of the ``/complete_weights_update`` POST and weight-version / + flush-cache propagation. +* The shape mismatch and "skip transfer" guards. + +Real Mooncake transfers are exercised by ``scripts/p2p_e2e_smoke.py`` +(needs GPUs + IB), not here. +""" + +from __future__ import annotations + +import socket +import sys +import types +from concurrent.futures import Future +from typing import Any, Dict, List, Tuple +from unittest.mock import patch + +import pytest +import torch + +from xorl.server.weight_sync.backends.base import EndpointConfig, TransportConfig +from xorl.server.weight_sync.backends.p2p import ( + P2PTransportBackend, + _async_api_enabled, + _resolve_local_hostname, + _transfer_small_entries, +) + + +class _FakeResponse: + def __init__(self, status_code: int = 200, payload: Dict[str, Any] | None = None): + self.status_code = status_code + self._payload = payload or {} + self.text = "" if status_code == 200 else "error body" + + def json(self) -> Dict[str, Any]: + return self._payload + + +class _FakeMooncakeEngine: + """Test double that records calls and returns success.""" + + def __init__(self, session_id: str = "fake-trainer:1234"): + self.session_id = session_id + self.engine = self + self.registered: List[Tuple[List[int], List[int]]] = [] + self.registered_memory: List[Tuple[int, int]] = [] + self.deregistered: List[List[int]] = [] + self.transfers: List[Tuple[str, List[int], List[int], List[int]]] = [] + self.fail_transfer = False + + # API surface mirrors MooncakeTransferEngine + def get_session_id(self) -> str: + return self.session_id + + def get_ib_device(self) -> str: + return "mlx5_0" + + def batch_register(self, ptrs: List[int], lengths: List[int]) -> int: + self.registered.append((list(ptrs), list(lengths))) + return 0 + + def register_memory(self, ptr: int, length: int) -> int: + self.registered_memory.append((ptr, length)) + return 0 + + def batch_deregister(self, ptrs: List[int]) -> int: + self.deregistered.append(list(ptrs)) + return 0 + + def batch_transfer_sync( + self, + session_id: str, + src_ptrs: List[int], + peer_ptrs: List[int], + lengths: List[int], + ) -> int: + self.transfers.append((session_id, list(src_ptrs), list(peer_ptrs), list(lengths))) + if self.fail_transfer: + return -1 + return 0 + + +def test_async_api_enabled_supports_warm_mode(monkeypatch): + monkeypatch.setenv("XORL_P2P_USE_ASYNC_API", "warm") + assert _async_api_enabled(cached_prepare=False) is False + assert _async_api_enabled(cached_prepare=True) is True + + +def test_transfer_small_entries_batches_by_session(monkeypatch): + monkeypatch.setenv("XORL_P2P_SMALL_TRANSFER_CHUNK", "2") + engine = _FakeMooncakeEngine() + session_bytes: Dict[str, int] = {} + session_transfer_s: Dict[str, float] = {} + + total_bytes, num_buffers = _transfer_small_entries( + engine_wrapper=engine, + small_session_data={ + "recv0:7000": [ + (0x1000, 0x2000, 4, None), + (0x1004, 0x2004, 4, None), + (0x1008, 0x2008, 4, None), + ] + }, + session_debug_info={"recv0:7000": {"session_id": "recv0:7000"}}, + small_register_ptrs=[0x1000], + small_register_lens=[12], + session_bytes=session_bytes, + session_transfer_s=session_transfer_s, + bucket_idx=3, + ) + + assert total_bytes == 12 + assert num_buffers == 3 + assert session_bytes == {"recv0:7000": 12} + assert "recv0:7000" in session_transfer_s + assert engine.registered == [([0x1000], [12])] + assert engine.deregistered == [[0x1000]] + assert engine.transfers == [ + ("recv0:7000", [0x1000, 0x1004], [0x2000, 0x2004], [4, 4]), + ("recv0:7000", [0x1008], [0x2008], [4]), + ] + + +def test_persistent_source_registration_skips_already_registered_ranges(): + backend, engine = _make_backend() + backend._intervals_per_cuda_segment = lambda intervals: list(intervals) # type: ignore[assignment] + backend._registered_intervals = [(0x3000, 0x3100)] + backend._registered_source_ptrs = [0x3000] + + backend._register_persistent_source_intervals( + [(0x3040, 0x3080), (0x5000, 0x5080)], + bucket_idx=7, + ) + + assert engine.registered == [([0x5000], [0x80])] + assert backend._registered_source_ptrs == [0x3000, 0x5000] + + backend._register_persistent_source_intervals( + [(0x5008, 0x5010)], + bucket_idx=8, + ) + + assert engine.registered == [([0x5000], [0x80])] + + +def _make_backend(num_endpoints: int = 1) -> Tuple[P2PTransportBackend, _FakeMooncakeEngine]: + cfg = TransportConfig( + endpoints=[EndpointConfig(host=f"infer-{i}", port=5000 + i, world_size=2) for i in range(num_endpoints)], + master_address="trainer-0", + master_port=0, + group_name="weight_sync_group", + buffer_size_mb=64, + device="cuda:0", + backend_config={"hostname": "trainer-0", "gpu_id": 0, "cpu_scratch_pool_bytes": 1024 * 1024}, + ) + backend = P2PTransportBackend(cfg) + fake_engine = _FakeMooncakeEngine() + # Pin the fake engine in place of the real Mooncake import path so that + # `initialize()` doesn't try to construct a TransferEngine. This also + # makes `transfer_bucket` / `destroy` see the fake. + backend._engine = fake_engine + backend._make_local_engine = lambda: fake_engine # type: ignore[assignment] + return backend, fake_engine + + +def _hf_locator( + *, + tp_rank: int, + full_shape: List[int], + slc: List[List[int]], + ptr: int, + nbytes: int, + session_id: str, + dtype: str = "bfloat16", +) -> Dict[str, Any]: + return { + "hf_name": "ignored-by-callers-of-this-helper", + "tp_rank": tp_rank, + "dp_rank": 0, + "ep_rank": -1, + "dtype": dtype, + "full_shape": full_shape, + "slice": slc, + "ptr": ptr, + "nbytes": nbytes, + "session_id": session_id, + } + + +class TestP2PInitializeHandshake: + def test_resolve_local_hostname_prefers_explicit_p2p_env(self, monkeypatch): + for env_name in ("P2P_TRAINER_HOSTNAME", "XORL_WEIGHT_SYNC_MASTER_ADDRESS", "POD_IP"): + monkeypatch.delenv(env_name, raising=False) + + monkeypatch.setenv("POD_IP", "10.0.0.3") + assert _resolve_local_hostname() == "10.0.0.3" + + monkeypatch.setenv("XORL_WEIGHT_SYNC_MASTER_ADDRESS", "10.0.0.2") + assert _resolve_local_hostname() == "10.0.0.2" + + monkeypatch.setenv("P2P_TRAINER_HOSTNAME", "10.0.0.1") + assert _resolve_local_hostname() == "10.0.0.1" + + def test_prepare_payload_uses_p2p_transport_and_engine_info(self): + backend, engine = _make_backend(num_endpoints=1) + + prepare_response = _FakeResponse( + 200, + { + "success": True, + "message": "ok", + "tensor_map": { + "model.layers.0.self_attn.q_proj.weight": [ + { + **_hf_locator( + tp_rank=0, + full_shape=[128, 64], + slc=[[0, 64], [0, 64]], + ptr=0xDEAD0000, + nbytes=64 * 64 * 2, + session_id="recv-0:7000", + ), + "hf_name": "model.layers.0.self_attn.q_proj.weight", + }, + ] + }, + "receiver_transfer_engine_infos": [ + {"tp_rank": 0, "session_id": "recv-0:7000"}, + ], + }, + ) + + with patch("requests.post", return_value=prepare_response) as posted: + ok = backend.initialize() + + assert ok is True + posted.assert_called_once() + url = posted.call_args.args[0] + body = posted.call_args.kwargs["json"] + assert url == "http://infer-0:5000/prepare_weights_update" + assert body["transport"] == "p2p" + assert body["sender_transfer_engine_info"]["session_id"] == engine.session_id + assert body["sender_transfer_engine_info"]["training_rank"] == 0 + assert body["group_name"] == "weight_sync_group" + assert "p2p_return_tensor_map" not in body + + def test_initialize_aggregates_tensor_map_across_endpoints(self, monkeypatch): + monkeypatch.setenv("XORL_P2P_PREPARE_WORKERS", "1") + backend, _ = _make_backend(num_endpoints=2) + + ep0_resp = _FakeResponse( + 200, + { + "success": True, + "message": "ok", + "tensor_map": { + "model.layers.0.self_attn.q_proj.weight": [ + { + **_hf_locator( + tp_rank=0, + full_shape=[128, 64], + slc=[[0, 64], [0, 64]], + ptr=0x1000, + nbytes=64 * 64 * 2, + session_id="recv0a:7000", + ), + "hf_name": "model.layers.0.self_attn.q_proj.weight", + }, + ], + }, + "receiver_transfer_engine_infos": [{"tp_rank": 0, "session_id": "recv0a:7000"}], + }, + ) + ep1_resp = _FakeResponse( + 200, + { + "success": True, + "message": "ok", + "tensor_map": { + "model.layers.0.self_attn.q_proj.weight": [ + { + **_hf_locator( + tp_rank=1, + full_shape=[128, 64], + slc=[[64, 128], [0, 64]], + ptr=0x2000, + nbytes=64 * 64 * 2, + session_id="recv1a:7000", + ), + "hf_name": "model.layers.0.self_attn.q_proj.weight", + }, + ], + }, + "receiver_transfer_engine_infos": [{"tp_rank": 0, "session_id": "recv1a:7000"}], + }, + ) + + with patch("requests.post", side_effect=[ep0_resp, ep1_resp]): + ok = backend.initialize() + + assert ok is True + locators = backend._tensor_map["model.layers.0.self_attn.q_proj.weight"] + assert len(locators) == 2 + assert {loc["endpoint_idx"] for loc in locators} == {0, 1} + assert sorted(loc["session_id"] for loc in locators) == ["recv0a:7000", "recv1a:7000"] + + def test_initialize_returns_false_on_http_error(self): + backend, _ = _make_backend() + with patch("requests.post", return_value=_FakeResponse(500, {})): + assert backend.initialize() is False + + def test_initialize_returns_false_when_remote_reports_failure(self): + backend, _ = _make_backend() + with patch( + "requests.post", + return_value=_FakeResponse(200, {"success": False, "message": "no engine"}), + ): + assert backend.initialize() is False + + def test_cached_prepare_can_reuse_existing_tensor_map(self): + backend, _ = _make_backend() + cached_map = { + "model.layers.0.self_attn.q_proj.weight": [ + { + **_hf_locator( + tp_rank=0, + full_shape=[128, 64], + slc=[[0, 64], [0, 64]], + ptr=0xDEAD0000, + nbytes=64 * 64 * 2, + session_id="recv-0:7000", + ), + "hf_name": "model.layers.0.self_attn.q_proj.weight", + }, + ] + } + backend._tensor_map = cached_map + backend._receiver_session_ids = ["recv-0:7000"] + + prepare_response = _FakeResponse( + 200, + { + "success": True, + "message": "cached", + "tensor_map": None, + "receiver_transfer_engine_infos": [ + {"tp_rank": 0, "session_id": "recv-0:7000"}, + ], + }, + ) + + with patch("requests.post", return_value=prepare_response) as posted: + ok = backend.initialize() + + assert ok is True + assert backend._tensor_map == cached_map + assert backend._tensor_map is cached_map + assert backend._last_prepare_returned_tensor_map is False + body = posted.call_args.kwargs["json"] + assert body["p2p_return_tensor_map"] is False + + def test_cached_prepare_retries_full_map_when_receiver_session_changes(self): + backend, _ = _make_backend() + param_name = "model.layers.0.self_attn.q_proj.weight" + backend._tensor_map = { + param_name: [ + { + **_hf_locator( + tp_rank=0, + full_shape=[128, 64], + slc=[[0, 64], [0, 64]], + ptr=0xDEAD0000, + nbytes=64 * 64 * 2, + session_id="recv-old:7000", + ), + "hf_name": param_name, + }, + ] + } + backend._receiver_session_ids = ["recv-old:7000"] + backend._prefer_cached_prepare = True + new_locator = { + **_hf_locator( + tp_rank=0, + full_shape=[128, 64], + slc=[[0, 64], [0, 64]], + ptr=0xBEEF0000, + nbytes=64 * 64 * 2, + session_id="recv-new:7000", + ), + "hf_name": param_name, + } + + cached_response = _FakeResponse( + 200, + { + "success": True, + "message": "cached", + "tensor_map": None, + "receiver_transfer_engine_infos": [{"tp_rank": 0, "session_id": "recv-new:7000"}], + }, + ) + full_response = _FakeResponse( + 200, + { + "success": True, + "message": "full", + "tensor_map": {param_name: [new_locator]}, + "receiver_transfer_engine_infos": [{"tp_rank": 0, "session_id": "recv-new:7000"}], + }, + ) + + with patch("requests.post", side_effect=[cached_response, full_response]) as posted: + ok = backend.initialize() + + assert ok is True + assert posted.call_count == 2 + assert posted.call_args_list[0].kwargs["json"]["p2p_return_tensor_map"] is False + assert "p2p_return_tensor_map" not in posted.call_args_list[1].kwargs["json"] + assert backend._receiver_session_ids == ["recv-new:7000"] + assert backend._tensor_map[param_name][0]["session_id"] == "recv-new:7000" + assert backend._tensor_map[param_name][0]["ptr"] == 0xBEEF0000 + assert backend._last_prepare_returned_tensor_map is True + + def test_cached_prepare_merges_partial_tensor_map_with_cached_endpoints(self, monkeypatch): + monkeypatch.setenv("XORL_P2P_PREPARE_WORKERS", "1") + backend, _ = _make_backend(num_endpoints=2) + name = "model.layers.0.self_attn.q_proj.weight" + cached_map = { + name: [ + { + **_hf_locator( + tp_rank=0, + full_shape=[128, 64], + slc=[[0, 64], [0, 64]], + ptr=0x1000, + nbytes=64 * 64 * 2, + session_id="recv0-old:7000", + ), + "hf_name": name, + "endpoint_idx": 0, + }, + { + **_hf_locator( + tp_rank=1, + full_shape=[128, 64], + slc=[[64, 128], [0, 64]], + ptr=0x2000, + nbytes=64 * 64 * 2, + session_id="recv1-old:7000", + ), + "hf_name": name, + "endpoint_idx": 1, + }, + ] + } + backend._tensor_map = cached_map + backend._receiver_session_ids = ["recv0-old:7000", "recv1-old:7000"] + backend._session_debug_info = { + "recv0-old:7000": {"session_id": "recv0-old:7000", "endpoint_idx": 0}, + "recv1-old:7000": {"session_id": "recv1-old:7000", "endpoint_idx": 1}, + } + backend._prefer_cached_prepare = True + + ep0_resp = _FakeResponse(200, {"success": True, "message": "cached", "tensor_map": None}) + ep1_resp = _FakeResponse( + 200, + { + "success": True, + "message": "full", + "tensor_map": { + name: [ + { + **_hf_locator( + tp_rank=1, + full_shape=[128, 64], + slc=[[64, 128], [0, 64]], + ptr=0x3000, + nbytes=64 * 64 * 2, + session_id="recv1-new:7000", + ), + "hf_name": name, + } + ] + }, + "receiver_transfer_engine_infos": [{"tp_rank": 0, "session_id": "recv1-new:7000"}], + }, + ) + + with patch("requests.post", side_effect=[ep0_resp, ep1_resp]): + ok = backend.initialize() + + assert ok is True + locators = backend._tensor_map[name] + assert [(loc["endpoint_idx"], loc["ptr"], loc["session_id"]) for loc in locators] == [ + (0, 0x1000, "recv0-old:7000"), + (1, 0x3000, "recv1-new:7000"), + ] + assert backend._receiver_session_ids == ["recv0-old:7000", "recv1-new:7000"] + assert backend._last_prepare_returned_tensor_map is True + + def test_cached_prepare_retries_full_map_when_receiver_rejects_flag(self): + backend, _ = _make_backend() + backend._tensor_map = { + "model.layers.0.self_attn.q_proj.weight": [ + { + **_hf_locator( + tp_rank=0, + full_shape=[128, 64], + slc=[[0, 64], [0, 64]], + ptr=0xDEAD0000, + nbytes=64 * 64 * 2, + session_id="recv-0:7000", + ), + "hf_name": "model.layers.0.self_attn.q_proj.weight", + }, + ] + } + backend._receiver_session_ids = ["recv-0:7000"] + full_response = _FakeResponse( + 200, + { + "success": True, + "message": "full", + "tensor_map": backend._tensor_map, + "receiver_transfer_engine_infos": [{"tp_rank": 0, "session_id": "recv-0:7000"}], + }, + ) + + with patch("requests.post", side_effect=[_FakeResponse(422, {}), full_response]) as posted: + ok = backend.initialize() + + assert ok is True + assert posted.call_count == 2 + assert posted.call_args_list[0].kwargs["json"]["p2p_return_tensor_map"] is False + assert "p2p_return_tensor_map" not in posted.call_args_list[1].kwargs["json"] + assert backend._last_prepare_returned_tensor_map is True + + def test_cached_prepare_retries_rejecting_endpoint_without_restarting_all(self, monkeypatch): + monkeypatch.setenv("XORL_P2P_PREPARE_WORKERS", "1") + backend, _ = _make_backend(num_endpoints=2) + param_name = "model.layers.0.self_attn.q_proj.weight" + backend._tensor_map = { + param_name: [ + { + **_hf_locator( + tp_rank=0, + full_shape=[128, 64], + slc=[[0, 64], [0, 64]], + ptr=0xDEAD0000, + nbytes=64 * 64 * 2, + session_id="recv-0:7000", + ), + "hf_name": param_name, + "endpoint_idx": 0, + }, + { + **_hf_locator( + tp_rank=1, + full_shape=[128, 64], + slc=[[64, 128], [0, 64]], + ptr=0xBEEF0000, + nbytes=64 * 64 * 2, + session_id="recv-1:7000", + ), + "hf_name": param_name, + "endpoint_idx": 1, + }, + ] + } + backend._receiver_session_ids = ["recv-0:7000", "recv-1:7000"] + backend._prefer_cached_prepare = True + ep1_full_locator = { + **_hf_locator( + tp_rank=1, + full_shape=[128, 64], + slc=[[64, 128], [0, 64]], + ptr=0x2000, + nbytes=64 * 64 * 2, + session_id="recv-1:7000", + ), + "hf_name": param_name, + } + ep0_cached_response = _FakeResponse( + 200, + { + "success": True, + "message": "cached", + "tensor_map": None, + "receiver_transfer_engine_infos": [{"tp_rank": 0, "session_id": "recv-0:7000"}], + }, + ) + ep1_full_response = _FakeResponse( + 200, + { + "success": True, + "message": "full", + "tensor_map": {param_name: [ep1_full_locator]}, + "receiver_transfer_engine_infos": [{"tp_rank": 0, "session_id": "recv-1:7000"}], + }, + ) + + with patch( + "requests.post", + side_effect=[ep0_cached_response, _FakeResponse(422, {}), ep1_full_response], + ) as posted: + ok = backend.initialize() + + assert ok is True + assert posted.call_count == 3 + assert [call.args[0] for call in posted.call_args_list] == [ + "http://infer-0:5000/prepare_weights_update", + "http://infer-1:5001/prepare_weights_update", + "http://infer-1:5001/prepare_weights_update", + ] + assert posted.call_args_list[0].kwargs["json"]["p2p_return_tensor_map"] is False + assert posted.call_args_list[1].kwargs["json"]["p2p_return_tensor_map"] is False + assert "p2p_return_tensor_map" not in posted.call_args_list[2].kwargs["json"] + locators = backend._tensor_map[param_name] + assert {loc["endpoint_idx"] for loc in locators} == {0, 1} + assert {loc["session_id"]: loc["ptr"] for loc in locators} == { + "recv-0:7000": 0xDEAD0000, + "recv-1:7000": 0x2000, + } + assert backend._receiver_session_ids == ["recv-0:7000", "recv-1:7000"] + assert backend._last_prepare_returned_tensor_map is True + + def test_complete_sync_preserves_cached_prepare_state(self): + backend, _ = _make_backend() + cached_map = { + "model.layers.0.self_attn.q_proj.weight": [ + { + **_hf_locator( + tp_rank=0, + full_shape=[128, 64], + slc=[[0, 64], [0, 64]], + ptr=0xDEAD0000, + nbytes=64 * 64 * 2, + session_id="recv-0:7000", + ), + "hf_name": "model.layers.0.self_attn.q_proj.weight", + }, + ] + } + backend._tensor_map = cached_map + backend._receiver_session_ids = ["recv-0:7000"] + backend._session_debug_info = {"recv-0:7000": {"session_id": "recv-0:7000"}} + + with patch("requests.post", return_value=_FakeResponse(200, {"success": True})): + backend.complete_sync() + + assert backend._tensor_map == cached_map + assert backend._receiver_session_ids == ["recv-0:7000"] + assert backend._session_debug_info == {"recv-0:7000": {"session_id": "recv-0:7000"}} + + def test_complete_sync_forwards_p2p_tied_weight_aliases(self): + backend, _ = _make_backend() + backend.config.backend_config["p2p_tied_weight_aliases"] = {"lm_head.weight": "model.embed_tokens.weight"} + + with patch("requests.post", return_value=_FakeResponse(200, {"success": True})) as posted: + backend.complete_sync() + + body = posted.call_args.kwargs["json"] + assert body["p2p_tied_weight_aliases"] == {"lm_head.weight": "model.embed_tokens.weight"} + + +class TestP2PSlicing: + def test_slice_source_for_locator_extracts_qkv_q_slice(self): + full = torch.arange(128 * 64, dtype=torch.bfloat16).reshape(128, 64) + loc = _hf_locator( + tp_rank=0, + full_shape=[128, 64], + slc=[[0, 64], [0, 64]], + ptr=0, + nbytes=64 * 64 * 2, + session_id="x", + ) + view = P2PTransportBackend._slice_source_for_locator("q_proj", full, loc) + assert view is not None + assert view.shape == (64, 64) + # Equal to the literal slice of the source. + assert torch.equal(view, full[0:64, 0:64]) + + def test_slice_source_for_locator_other_rank_picks_other_rows(self): + full = torch.arange(128 * 64, dtype=torch.bfloat16).reshape(128, 64) + loc = _hf_locator( + tp_rank=1, + full_shape=[128, 64], + slc=[[64, 128], [0, 64]], + ptr=0, + nbytes=64 * 64 * 2, + session_id="x", + ) + view = P2PTransportBackend._slice_source_for_locator("q_proj", full, loc) + assert view is not None + assert torch.equal(view, full[64:128, 0:64]) + + def test_slice_source_returns_none_on_full_shape_mismatch(self): + full = torch.zeros(128, 64, dtype=torch.bfloat16) + loc = _hf_locator( + tp_rank=0, + full_shape=[256, 64], # wrong on purpose + slc=[[0, 128], [0, 64]], + ptr=0, + nbytes=128 * 64 * 2, + session_id="x", + ) + assert P2PTransportBackend._slice_source_for_locator("q_proj", full, loc) is None + + def test_slice_source_no_slice_returns_full_tensor(self): + full = torch.zeros(8, 16, dtype=torch.bfloat16) + loc = {"ptr": 0, "nbytes": 8 * 16 * 2} + view = P2PTransportBackend._slice_source_for_locator("rn", full, loc) + assert view is full + + def test_slice_source_squeezes_qwen35_linear_attention_conv_for_receiver_layout(self): + full = torch.arange(8 * 1 * 4, dtype=torch.bfloat16).reshape(8, 1, 4) + loc = _hf_locator( + tp_rank=1, + full_shape=[8, 4], + slc=[[4, 8], [0, 4]], + ptr=0, + nbytes=4 * 4 * 2, + session_id="x", + ) + view = P2PTransportBackend._slice_source_for_locator( + "model.layers.0.linear_attn.conv1d.weight", + full, + loc, + ) + assert view is not None + assert view.shape == (4, 4) + assert torch.equal(view, full.squeeze(1)[4:8, 0:4]) + + def test_slice_source_handles_qwen35_linear_attention_local_state_vector(self): + full = torch.arange(32, dtype=torch.float32) + loc = _hf_locator( + tp_rank=2, + full_shape=[8], + slc=[[0, 8]], + ptr=0, + nbytes=8 * 4, + session_id="x", + dtype="float32", + ) + view = P2PTransportBackend._slice_source_for_locator( + "model.layers.0.linear_attn.A_log", + full, + loc, + ) + assert view is not None + assert view.shape == (8,) + assert torch.equal(view, full[16:24]) + + def test_slice_source_casts_qwen35_linear_attention_state_to_receiver_dtype(self): + full = torch.arange(32, dtype=torch.bfloat16) + loc = _hf_locator( + tp_rank=2, + full_shape=[32], + slc=[[16, 24]], + ptr=0, + nbytes=8 * 4, + session_id="x", + dtype="float32", + ) + view = P2PTransportBackend._slice_source_for_locator( + "model.layers.0.linear_attn.A_log", + full, + loc, + ) + assert view is not None + assert view.shape == (8,) + assert view.dtype == torch.float32 + assert view.numel() * view.element_size() == 8 * 4 + assert torch.equal(view, full[16:24].float()) + + +class TestP2PTransferBucket: + def _seed_tensor_map(self, backend: P2PTransportBackend, peer_ptr: int): + backend._tensor_map = { + "model.layers.0.self_attn.q_proj.weight": [ + { + **_hf_locator( + tp_rank=0, + full_shape=[128, 64], + slc=[[0, 64], [0, 64]], + ptr=peer_ptr, + nbytes=64 * 64 * 2, + session_id="recv0:7000", + ), + "hf_name": "model.layers.0.self_attn.q_proj.weight", + "endpoint_idx": 0, + }, + { + **_hf_locator( + tp_rank=1, + full_shape=[128, 64], + slc=[[64, 128], [0, 64]], + ptr=peer_ptr + 0x10000, + nbytes=64 * 64 * 2, + session_id="recv1:7000", + ), + "hf_name": "model.layers.0.self_attn.q_proj.weight", + "endpoint_idx": 0, + }, + ] + } + backend._receiver_session_ids = ["recv0:7000", "recv1:7000"] + + def test_transfer_bucket_writes_correct_slice_per_receiver(self): + backend, engine = _make_backend() + self._seed_tensor_map(backend, peer_ptr=0xCAFE_0000) + full = torch.arange(128 * 64, dtype=torch.bfloat16).reshape(128, 64) + + backend.transfer_bucket( + [("model.layers.0.self_attn.q_proj.weight", full)], + src_rank=0, + ) + backend.flush_pending_transfers() + + # Two transfers issued — one per receiver. + sessions = sorted(t[0] for t in engine.transfers) + assert sessions == ["recv0:7000", "recv1:7000"] + + # Per-session: one buffer of 64*64*2 bytes. + for session_id, src_ptrs, peer_ptrs, lengths in engine.transfers: + assert len(src_ptrs) == len(peer_ptrs) == len(lengths) == 1 + assert lengths[0] == 64 * 64 * 2 + if session_id == "recv0:7000": + assert peer_ptrs[0] == 0xCAFE_0000 + else: + assert peer_ptrs[0] == 0xCAFE_0000 + 0x10000 + + def test_transfer_bucket_reuses_staged_source_for_replicated_locators(self): + backend, engine = _make_backend(num_endpoints=2) + backend._tensor_map = { + "param": [ + { + **_hf_locator( + tp_rank=0, + full_shape=[8, 8], + slc=[[0, 8], [0, 8]], + ptr=0x1000, + nbytes=8 * 8 * 2, + session_id="recv0:7000", + ), + "hf_name": "param", + "endpoint_idx": 0, + }, + { + **_hf_locator( + tp_rank=0, + full_shape=[8, 8], + slc=[[0, 8], [0, 8]], + ptr=0x2000, + nbytes=8 * 8 * 2, + session_id="recv1:7000", + ), + "hf_name": "param", + "endpoint_idx": 1, + }, + ] + } + backend._receiver_session_ids = ["recv0:7000", "recv1:7000"] + full = torch.arange(8 * 8, dtype=torch.bfloat16).reshape(8, 8) + + backend.transfer_bucket([("param", full)], src_rank=0) + backend.flush_pending_transfers() + + assert len(engine.transfers) == 2 + assert engine.transfers[0][1] == engine.transfers[1][1] + assert [row[2] for row in engine.transfers] == [[0x1000], [0x2000]] + assert [row[3] for row in engine.transfers] == [[8 * 8 * 2], [8 * 8 * 2]] + + def test_transfer_bucket_fails_on_unknown_param(self): + backend, engine = _make_backend() + backend._tensor_map = {} + full = torch.zeros(8, 8, dtype=torch.bfloat16) + with pytest.raises(RuntimeError, match="no receiver locator"): + backend.transfer_bucket([("unknown.param", full)], src_rank=0) + assert engine.transfers == [] + + def test_transfer_bucket_skips_missing_tied_lm_head_locator(self): + backend, engine = _make_backend() + receiver_name = "model.embed_tokens.weight" + backend._tensor_map = { + receiver_name: [ + { + **_hf_locator( + tp_rank=0, + full_shape=[8, 8], + slc=[[0, 8], [0, 8]], + ptr=0x1234_0000, + nbytes=8 * 8 * 2, + session_id="recv0:7000", + ), + "hf_name": receiver_name, + "endpoint_idx": 0, + } + ] + } + backend._receiver_session_ids = ["recv0:7000"] + full = torch.zeros(8, 8, dtype=torch.bfloat16) + + backend.transfer_bucket( + [ + ("model.embed_tokens.weight", full), + ("lm_head.weight", full.clone()), + ], + src_rank=0, + ) + backend.flush_pending_transfers() + + assert len(engine.transfers) == 1 + session_id, _src_ptrs, peer_ptrs, lengths = engine.transfers[0] + assert session_id == "recv0:7000" + assert peer_ptrs == [0x1234_0000] + assert lengths == [full.numel() * full.element_size()] + + def test_transfer_bucket_uses_language_model_receiver_prefix_fallback(self): + backend, engine = _make_backend() + receiver_name = "language_model.model.layers.0.self_attn.q_b_proj.weight" + backend._tensor_map = { + receiver_name: [ + { + **_hf_locator( + tp_rank=0, + full_shape=[4, 8], + slc=[[0, 4], [0, 8]], + ptr=0x4567_0000, + nbytes=4 * 8 * 2, + session_id="recv0:7000", + ), + "hf_name": receiver_name, + "endpoint_idx": 0, + } + ] + } + backend._receiver_session_ids = ["recv0:7000"] + full = torch.zeros(4, 8, dtype=torch.bfloat16) + + backend.transfer_bucket( + [("model.layers.0.self_attn.q_b_proj.weight", full)], + src_rank=0, + ) + backend.flush_pending_transfers() + + assert len(engine.transfers) == 1 + session_id, _src_ptrs, peer_ptrs, lengths = engine.transfers[0] + assert session_id == "recv0:7000" + assert peer_ptrs == [0x4567_0000] + assert lengths == [full.numel() * full.element_size()] + + def test_transfer_bucket_fails_on_shape_mismatch(self): + backend, engine = _make_backend() + backend._tensor_map = { + "param": [ + _hf_locator( + tp_rank=0, + full_shape=[16, 8], + slc=[[0, 16], [0, 8]], + ptr=0x1000, + nbytes=16 * 8 * 2, + session_id="recv0:7000", + ) + ] + } + full = torch.zeros(8, 8, dtype=torch.bfloat16) + + with pytest.raises(RuntimeError, match="incomplete or incompatible"): + backend.transfer_bucket([("param", full)], src_rank=0) + + assert engine.transfers == [] + + def test_transfer_bucket_rejects_nonzero_src_rank(self): + backend, _ = _make_backend() + backend._tensor_map = {} + try: + backend.transfer_bucket([], src_rank=2) + except ValueError as e: + assert "src_rank=0" in str(e) or "src_rank" in str(e) + else: + raise AssertionError("expected ValueError for non-zero src_rank") + + def test_transfer_bucket_stashes_weight_version_and_flush_for_destroy(self): + backend, _ = _make_backend() + self._seed_tensor_map(backend, peer_ptr=0x1000) + full = torch.zeros(128, 64, dtype=torch.bfloat16) + backend.transfer_bucket( + [("model.layers.0.self_attn.q_proj.weight", full)], + src_rank=0, + flush_cache=True, + weight_version="rev-42", + ) + backend.flush_pending_transfers() + assert backend.config.backend_config["flush_cache"] is True + assert backend.config.backend_config["weight_version"] == "rev-42" + + def test_transfer_bucket_preserves_requested_flush_cache(self): + backend, _ = _make_backend() + self._seed_tensor_map(backend, peer_ptr=0x1000) + backend.config.backend_config["flush_cache"] = True + full = torch.zeros(128, 64, dtype=torch.bfloat16) + + backend.transfer_bucket( + [("model.layers.0.self_attn.q_proj.weight", full)], + src_rank=0, + flush_cache=False, + ) + + with patch( + "requests.post", + return_value=_FakeResponse(200, {"success": True, "message": "ok"}), + ) as posted: + backend.complete_sync() + + body = posted.call_args.kwargs["json"] + assert body["flush_cache"] is True + + def test_transfer_bucket_stages_small_cpu_tensors_through_cpu_pool(self): + backend, engine = _make_backend() + backend._tensor_map = { + "model.layers.0.mlp.gate_proj.weight_scale_inv": [ + { + **_hf_locator( + tp_rank=0, + full_shape=[2, 2], + slc=[[0, 2], [0, 2]], + ptr=0x1234_0000, + nbytes=2 * 2 * 4, + session_id="recv0:7000", + dtype="float32", + ), + "hf_name": "model.layers.0.mlp.gate_proj.weight_scale_inv", + "endpoint_idx": 0, + } + ] + } + backend._receiver_session_ids = ["recv0:7000"] + scale = torch.ones(2, 2, dtype=torch.float32) + + backend.transfer_bucket( + [("model.layers.0.mlp.gate_proj.weight_scale_inv", scale)], + src_rank=0, + ) + backend.flush_pending_transfers() + + assert engine.registered_memory + assert engine.registered == [] + assert len(engine.transfers) == 1 + session_id, _src_ptrs, peer_ptrs, lengths = engine.transfers[0] + assert session_id == "recv0:7000" + assert peer_ptrs == [0x1234_0000] + assert lengths == [scale.numel() * scale.element_size()] + + def test_transfer_bucket_aligns_cpu_scratch_offsets_before_dtype_views(self, monkeypatch): + monkeypatch.setenv("XORL_P2P_MOONCAKE_TRANSFER_CHUNK", "16") + backend, engine = _make_backend() + backend._tensor_map = { + "model.layers.0.mlp.gate_proj.weight": [ + { + **_hf_locator( + tp_rank=0, + full_shape=[3], + slc=[[0, 3]], + ptr=0x2000, + nbytes=3, + session_id="recv0:7000", + dtype="float8_e4m3fn", + ), + "hf_name": "model.layers.0.mlp.gate_proj.weight", + } + ], + "model.layers.0.mlp.gate_proj.weight_scale_inv": [ + { + **_hf_locator( + tp_rank=0, + full_shape=[1], + slc=[[0, 1]], + ptr=0x2003, + nbytes=4, + session_id="recv0:7000", + dtype="float32", + ), + "hf_name": "model.layers.0.mlp.gate_proj.weight_scale_inv", + } + ], + } + backend._receiver_session_ids = ["recv0:7000"] + + backend.transfer_bucket( + [ + ("model.layers.0.mlp.gate_proj.weight", torch.ones(3, dtype=torch.uint8)), + ("model.layers.0.mlp.gate_proj.weight_scale_inv", torch.ones(1, dtype=torch.float32)), + ], + src_rank=0, + ) + backend.flush_pending_transfers() + + assert len(engine.transfers) == 1 + session_id, src_ptrs, peer_ptrs, lengths = engine.transfers[0] + assert session_id == "recv0:7000" + assert peer_ptrs == [0x2000, 0x2003] + assert lengths == [3, 4] + assert src_ptrs[1] % 4 == 0 + + def test_transfer_bucket_does_not_coalesce_across_receiver_memory_handles(self, monkeypatch): + monkeypatch.setenv("XORL_P2P_MOONCAKE_TRANSFER_CHUNK", "16") + backend, engine = _make_backend() + backend._tensor_map = { + "param_a": [ + { + **_hf_locator( + tp_rank=0, + full_shape=[2], + slc=[[0, 2]], + ptr=0x1000, + nbytes=4, + session_id="recv0:7000", + ), + "hf_name": "param_a", + "memory_handle": 0x1000, + } + ], + "param_b": [ + { + **_hf_locator( + tp_rank=0, + full_shape=[2], + slc=[[0, 2]], + ptr=0x1004, + nbytes=4, + session_id="recv0:7000", + ), + "hf_name": "param_b", + "memory_handle": 0x2000, + } + ], + } + backend._receiver_session_ids = ["recv0:7000"] + + backend.transfer_bucket( + [ + ("param_a", torch.ones(2, dtype=torch.bfloat16)), + ("param_b", torch.ones(2, dtype=torch.bfloat16)), + ], + src_rank=0, + ) + backend.flush_pending_transfers() + + assert len(engine.transfers) == 1 + session_id, _src_ptrs, peer_ptrs, lengths = engine.transfers[0] + assert session_id == "recv0:7000" + assert peer_ptrs == [0x1000, 0x1004] + assert lengths == [4, 4] + + def test_transfer_bucket_does_not_coalesce_without_receiver_memory_handles(self, monkeypatch): + monkeypatch.setenv("XORL_P2P_MOONCAKE_TRANSFER_CHUNK", "16") + backend, engine = _make_backend() + backend._tensor_map = { + "param_a": [ + { + **_hf_locator( + tp_rank=0, + full_shape=[2], + slc=[[0, 2]], + ptr=0x1000, + nbytes=4, + session_id="recv0:7000", + ), + "hf_name": "param_a", + } + ], + "param_b": [ + { + **_hf_locator( + tp_rank=0, + full_shape=[2], + slc=[[0, 2]], + ptr=0x1004, + nbytes=4, + session_id="recv0:7000", + ), + "hf_name": "param_b", + } + ], + } + backend._receiver_session_ids = ["recv0:7000"] + + backend.transfer_bucket( + [ + ("param_a", torch.ones(2, dtype=torch.bfloat16)), + ("param_b", torch.ones(2, dtype=torch.bfloat16)), + ], + src_rank=0, + ) + backend.flush_pending_transfers() + + assert len(engine.transfers) == 1 + session_id, _src_ptrs, peer_ptrs, lengths = engine.transfers[0] + assert session_id == "recv0:7000" + assert peer_ptrs == [0x1000, 0x1004] + assert lengths == [4, 4] + + def test_transfer_bucket_failure_names_tensor_and_receiver_handle(self, monkeypatch): + monkeypatch.setenv("XORL_P2P_TRANSFER_RETRIES", "1") + monkeypatch.setenv("XORL_P2P_TRANSFER_DEBUG", "1") + backend, engine = _make_backend() + engine.fail_transfer = True + backend._tensor_map = { + "model.layers.1.mlp.experts.0.gate_proj.weight_scale_inv": [ + { + **_hf_locator( + tp_rank=0, + full_shape=[2], + slc=[[0, 2]], + ptr=0x3000, + nbytes=8, + session_id="recv0:7000", + dtype="float32", + ), + "hf_name": "model.layers.1.mlp.experts.0.gate_proj.weight_scale_inv", + "memory_handle": 0x3000, + } + ] + } + backend._receiver_session_ids = ["recv0:7000"] + + backend.transfer_bucket( + [ + ( + "model.layers.1.mlp.experts.0.gate_proj.weight_scale_inv", + torch.ones(2, dtype=torch.float32), + ) + ], + src_rank=0, + ) + with pytest.raises(RuntimeError) as exc_info: + backend.flush_pending_transfers() + + message = str(exc_info.value) + assert "weight_scale_inv" in message + assert "handle=0x3000" in message + assert "ptr=0x3000" in message + + def test_transfer_bucket_failure_caps_coalesced_debug_sample(self, monkeypatch): + monkeypatch.setenv("XORL_P2P_TRANSFER_RETRIES", "1") + monkeypatch.setenv("XORL_P2P_TRANSFER_DEBUG", "1") + backend, engine = _make_backend() + engine.fail_transfer = True + locators = [] + for idx in range(8): + locators.append( + { + **_hf_locator( + tp_rank=0, + full_shape=[8], + slc=[[idx, idx + 1]], + ptr=0x4000 + idx * 2, + nbytes=2, + session_id="recv0:7000", + ), + "hf_name": "param", + "memory_handle": 0x4000, + } + ) + backend._tensor_map = {"param": locators} + backend._receiver_session_ids = ["recv0:7000"] + + backend.transfer_bucket([("param", torch.ones(8, dtype=torch.bfloat16))], src_rank=0) + with pytest.raises(RuntimeError) as exc_info: + backend.flush_pending_transfers() + + message = str(exc_info.value) + assert "ptr=0x4000" in message + assert "ptr=0x400a" in message + assert "... 2 more" in message + + def test_transfer_bucket_failure_skips_debug_samples_by_default(self, monkeypatch): + monkeypatch.setenv("XORL_P2P_TRANSFER_RETRIES", "1") + monkeypatch.delenv("XORL_P2P_TRANSFER_DEBUG", raising=False) + backend, engine = _make_backend() + engine.fail_transfer = True + backend._tensor_map = { + "param": [ + { + **_hf_locator( + tp_rank=0, + full_shape=[1], + slc=[[0, 1]], + ptr=0x5000, + nbytes=2, + session_id="recv0:7000", + ), + "hf_name": "param", + "memory_handle": 0x5000, + } + ] + } + backend._receiver_session_ids = ["recv0:7000"] + + backend.transfer_bucket([("param", torch.ones(1, dtype=torch.bfloat16))], src_rank=0) + with pytest.raises(RuntimeError) as exc_info: + backend.flush_pending_transfers() + + message = str(exc_info.value) + assert "transfer_debug=disabled" in message + assert "ptr=0x5000" not in message + + +class TestP2PDirectEPTransfer: + def _make_multi_sender_backend( + self, + *, + rank_index: int, + world_size: int, + rank_filter, + sender_ranks=None, + process_group=None, + direct_ep_size=None, + sender_ep_ranks=None, + direct_ep_dense_sharding=False, + ) -> Tuple[P2PTransportBackend, _FakeMooncakeEngine]: + backend_config = { + "hostname": "trainer-0", + "gpu_id": 0, + "direct_ep_transfer": True, + "world_size": world_size, + "rank_index": rank_index, + "rank_filter": rank_filter, + "cpu_scratch_pool_bytes": 1024 * 1024, + "direct_ep_dense_sharding": direct_ep_dense_sharding, + } + if sender_ranks is not None: + backend_config["sender_ranks"] = sender_ranks + if process_group is not None: + backend_config["process_group"] = process_group + if direct_ep_size is not None: + backend_config["direct_ep_size"] = direct_ep_size + if sender_ep_ranks is not None: + backend_config["sender_ep_ranks"] = sender_ep_ranks + cfg = TransportConfig( + endpoints=[EndpointConfig(host="infer-0", port=5000, world_size=2)], + master_address="trainer-0", + master_port=0, + group_name="weight_sync_group", + buffer_size_mb=64, + device="cuda:0", + training_rank=rank_index, + backend_config=backend_config, + ) + backend = P2PTransportBackend(cfg) + fake_engine = _FakeMooncakeEngine(session_id=f"trainer-{rank_index}:1234") + backend._engine = fake_engine + backend._make_local_engine = lambda: fake_engine # type: ignore[assignment] + return backend, fake_engine + + def test_supports_direct_ep_transfer_advertises_all_ranks(self): + backend, _ = self._make_multi_sender_backend(rank_index=0, world_size=4, rank_filter=lambda loc: True) + assert backend.supports_direct_ep_transfer is True + assert backend.supports_direct_pp_transfer is False + assert backend.sender_ranks == frozenset({0, 1, 2, 3}) + assert backend.has_explicit_sender_ranks is False + + def test_direct_ep_transfer_accepts_explicit_sender_ranks(self): + backend, _ = self._make_multi_sender_backend( + rank_index=0, + world_size=4, + rank_filter=lambda loc: True, + sender_ranks=(0, 2), + ) + assert backend.sender_ranks == frozenset({0, 2}) + assert backend.has_explicit_sender_ranks is True + + def test_initialize_multi_sender_scatters_filtered_tensor_maps(self, monkeypatch): + monkeypatch.delenv("XORL_P2P_SCATTER_REUSE_LOCATORS", raising=False) + monkeypatch.delenv("XORL_P2P_SCATTER_COPY_MODE", raising=False) + process_group = object() + backend, _ = self._make_multi_sender_backend( + rank_index=0, + world_size=4, + rank_filter=lambda loc: True, + sender_ranks=(0, 1, 2, 3), + process_group=process_group, + direct_ep_size=4, + sender_ep_ranks=((0, 0), (1, 1), (2, 2), (3, 3)), + ) + dense_name = "model.embed_tokens.weight" + expert_names = [f"model.layers.0.mlp.experts.{idx}.gate_proj.weight" for idx in range(4)] + tensor_map = { + dense_name: [ + { + **_hf_locator( + tp_rank=0, + full_shape=[8, 8], + slc=[[0, 8], [0, 8]], + ptr=0x1000, + nbytes=8 * 8 * 2, + session_id="recv:7000", + ), + "hf_name": dense_name, + "endpoint_idx": 0, + } + ], + **{ + name: [ + { + **_hf_locator( + tp_rank=0, + full_shape=[8, 8], + slc=[[0, 8], [0, 8]], + ptr=0x2000 + idx * 0x100, + nbytes=8 * 8 * 2, + session_id="recv:7000", + ), + "hf_name": name, + "endpoint_idx": 0, + } + ] + for idx, name in enumerate(expert_names) + }, + } + + def initialize_rank0(): + backend._tensor_map = tensor_map + backend._receiver_session_ids = ["recv:7000"] + backend._session_debug_info = {"recv:7000": {"session_id": "recv:7000", "endpoint_idx": 0}} + backend._last_prepare_returned_tensor_map = True + backend._last_prepare_tensor_map_endpoint_indices = {0} + return True + + all_gather_calls = [] + scatter_inputs = [] + + def all_gather_object(output, value, **kwargs): + assert kwargs.get("group") is process_group + all_gather_calls.append(value) + output[:] = [False, False, False, False] if len(all_gather_calls) == 1 else [True] * 4 + + def scatter_object_list(output, scatter_input, **kwargs): + assert kwargs.get("group") is process_group + assert kwargs.get("src") == 0 + scatter_inputs.append(scatter_input) + output[0] = scatter_input[0] + + backend._initialize_single_sender = initialize_rank0 # type: ignore[assignment] + + with patch("xorl.server.weight_sync.backends.p2p.dist.is_available", return_value=True): + with patch("xorl.server.weight_sync.backends.p2p.dist.is_initialized", return_value=True): + with patch("xorl.server.weight_sync.backends.p2p.dist.broadcast_object_list") as broadcast: + with patch( + "xorl.server.weight_sync.backends.p2p.dist.scatter_object_list", + side_effect=scatter_object_list, + ): + with patch( + "xorl.server.weight_sync.backends.p2p.dist.all_gather_object", + side_effect=all_gather_object, + ): + assert backend.initialize() is True + + broadcast.assert_not_called() + payloads = scatter_inputs[0] + assert payloads[0] == ("rank0_ready",) + assert payloads[1][0] == "tensor_map_with_infos" + assert set(payloads[1][1]) == {expert_names[1]} + assert payloads[1][1][expert_names[1]] == tensor_map[expert_names[1]] + assert payloads[1][1][expert_names[1]] is tensor_map[expert_names[1]] + assert payloads[1][1][expert_names[1]][0] is tensor_map[expert_names[1]][0] + assert payloads[2][0] == "tensor_map_with_infos" + assert set(payloads[2][1]) == {expert_names[2]} + assert dense_name not in payloads[1][1] + + def test_scatter_copy_mode_list_copies_locator_lists(self, monkeypatch): + monkeypatch.setenv("XORL_P2P_SCATTER_COPY_MODE", "list") + monkeypatch.delenv("XORL_P2P_SCATTER_REUSE_LOCATORS", raising=False) + backend, _ = self._make_multi_sender_backend( + rank_index=0, + world_size=2, + rank_filter=lambda loc: True, + sender_ranks=(0, 1), + direct_ep_size=2, + sender_ep_ranks=((0, 0), (1, 1)), + ) + expert_name = "model.layers.0.mlp.experts.1.gate_proj.weight" + tensor_map = { + expert_name: [ + { + **_hf_locator( + tp_rank=0, + full_shape=[8, 8], + slc=[[0, 8], [0, 8]], + ptr=0x2000, + nbytes=8 * 8 * 2, + session_id="recv:7000", + ), + "hf_name": expert_name, + } + ] + } + + filtered = backend._filter_tensor_map_for_sender( + tensor_map, + 1, + locator_copy_mode=backend._scatter_locator_copy_mode(), + ) + + assert filtered[expert_name] == tensor_map[expert_name] + assert filtered[expert_name] is not tensor_map[expert_name] + assert filtered[expert_name][0] is tensor_map[expert_name][0] + + def test_scatter_reuse_flag_overrides_copy_mode_for_existing_manifests(self, monkeypatch): + monkeypatch.setenv("XORL_P2P_SCATTER_COPY_MODE", "list") + monkeypatch.setenv("XORL_P2P_SCATTER_REUSE_LOCATORS", "1") + + assert P2PTransportBackend._scatter_locator_copy_mode() == "none" + + def test_scatter_copy_mode_deep_copies_locator_dicts(self, monkeypatch): + monkeypatch.setenv("XORL_P2P_SCATTER_COPY_MODE", "deep") + monkeypatch.delenv("XORL_P2P_SCATTER_REUSE_LOCATORS", raising=False) + backend, _ = self._make_multi_sender_backend( + rank_index=0, + world_size=2, + rank_filter=lambda loc: True, + sender_ranks=(0, 1), + direct_ep_size=2, + sender_ep_ranks=((0, 0), (1, 1)), + ) + expert_name = "model.layers.0.mlp.experts.1.gate_proj.weight" + tensor_map = { + expert_name: [ + { + **_hf_locator( + tp_rank=0, + full_shape=[8, 8], + slc=[[0, 8], [0, 8]], + ptr=0x2000, + nbytes=8 * 8 * 2, + session_id="recv:7000", + ), + "hf_name": expert_name, + } + ] + } + + filtered = backend._filter_tensor_map_for_sender( + tensor_map, + 1, + locator_copy_mode=backend._scatter_locator_copy_mode(), + ) + + assert filtered[expert_name] == tensor_map[expert_name] + assert filtered[expert_name] is not tensor_map[expert_name] + assert filtered[expert_name][0] is not tensor_map[expert_name][0] + + def test_direct_ep_dense_sharding_partitions_dense_tensor_maps(self): + backend, _ = self._make_multi_sender_backend( + rank_index=0, + world_size=4, + rank_filter=lambda loc: True, + sender_ranks=(0, 1, 2, 3), + direct_ep_size=4, + sender_ep_ranks=((0, 0), (1, 1), (2, 2), (3, 3)), + direct_ep_dense_sharding=True, + ) + dense_names = [ + "model.embed_tokens.weight", + "model.layers.0.self_attn.qkv_proj.weight", + "model.layers.0.self_attn.q_proj.weight", + "model.layers.0.self_attn.k_proj.weight", + "model.layers.0.self_attn.v_proj.weight", + "model.layers.0.mlp.gate_up_proj.weight", + "model.layers.0.mlp.gate_proj.weight", + "model.layers.0.mlp.up_proj.weight", + "model.layers.0.self_attn.o_proj.weight", + "model.norm.weight", + ] + expert_names = [f"model.layers.0.mlp.experts.{idx}.gate_proj.weight" for idx in range(4)] + tensor_map = { + name: [ + { + **_hf_locator( + tp_rank=0, + full_shape=[8, 8], + slc=[[0, 8], [0, 8]], + ptr=0x1000, + nbytes=8 * 8 * 2, + session_id="recv:7000", + ), + "hf_name": name, + } + ] + for name in dense_names + expert_names + } + + owners = { + name: [rank for rank in backend.sender_rank_order if backend.should_send_dense_param(name, rank)] + for name in dense_names + } + assert all(len(ranks) == 1 for ranks in owners.values()) + assert owners["model.layers.0.self_attn.qkv_proj.weight"] == owners["model.layers.0.self_attn.q_proj.weight"] + assert owners["model.layers.0.self_attn.qkv_proj.weight"] == owners["model.layers.0.self_attn.k_proj.weight"] + assert owners["model.layers.0.self_attn.qkv_proj.weight"] == owners["model.layers.0.self_attn.v_proj.weight"] + assert owners["model.layers.0.mlp.gate_up_proj.weight"] == owners["model.layers.0.mlp.gate_proj.weight"] + assert owners["model.layers.0.mlp.gate_up_proj.weight"] == owners["model.layers.0.mlp.up_proj.weight"] + + for sender_rank in backend.sender_rank_order: + filtered = backend._filter_tensor_map_for_sender(tensor_map, sender_rank) + expected_dense = {name for name, ranks in owners.items() if ranks == [sender_rank]} + expected_expert = {expert_names[sender_rank]} + assert set(filtered) == expected_dense | expected_expert + + def test_direct_ep_dense_sharding_filters_dense_buffers_by_sender(self): + backend, _ = self._make_multi_sender_backend( + rank_index=0, + world_size=4, + rank_filter=lambda loc: True, + sender_ranks=(0, 1, 2, 3), + direct_ep_dense_sharding=True, + ) + buffer = [ + ("model.embed_tokens.weight", torch.ones(1)), + ("model.layers.0.self_attn.qkv_proj.weight", torch.ones(1)), + ("model.layers.0.self_attn.q_proj.weight", torch.ones(1)), + ("model.layers.0.self_attn.k_proj.weight", torch.ones(1)), + ("model.layers.0.self_attn.v_proj.weight", torch.ones(1)), + ("model.layers.0.mlp.experts.0.gate_proj.weight", torch.ones(1)), + ] + + for sender_rank in backend.sender_rank_order: + filtered_names = {name for name, _ in backend.filter_dense_buffer_for_rank(buffer, sender_rank)} + assert "model.layers.0.mlp.experts.0.gate_proj.weight" not in filtered_names + assert filtered_names == {name for name, _ in buffer if backend.should_send_dense_param(name, sender_rank)} + + def test_initialize_multi_sender_nonzero_adopts_scattered_tensor_map(self): + process_group = object() + backend, _ = self._make_multi_sender_backend( + rank_index=2, + world_size=4, + rank_filter=lambda loc: True, + sender_ranks=(0, 1, 2, 3), + process_group=process_group, + direct_ep_size=4, + sender_ep_ranks=((0, 0), (1, 1), (2, 2), (3, 3)), + ) + expert_name = "model.layers.0.mlp.experts.2.gate_proj.weight" + payload = ( + "tensor_map_with_infos", + { + expert_name: [ + { + **_hf_locator( + tp_rank=0, + full_shape=[8, 8], + slc=[[0, 8], [0, 8]], + ptr=0x2200, + nbytes=8 * 8 * 2, + session_id="recv:7000", + ), + "hf_name": expert_name, + "endpoint_idx": 0, + } + ] + }, + [{"session_id": "recv:7000", "endpoint_idx": 0}], + ) + all_gather_calls = [] + + def all_gather_object(output, value, **kwargs): + all_gather_calls.append(value) + output[:] = [False, False, False, False] if len(all_gather_calls) == 1 else [True] * 4 + + def scatter_object_list(output, scatter_input, **kwargs): + assert scatter_input is None + output[0] = payload + + with patch("xorl.server.weight_sync.backends.p2p.dist.is_available", return_value=True): + with patch("xorl.server.weight_sync.backends.p2p.dist.is_initialized", return_value=True): + with patch( + "xorl.server.weight_sync.backends.p2p.dist.scatter_object_list", + side_effect=scatter_object_list, + ): + with patch( + "xorl.server.weight_sync.backends.p2p.dist.all_gather_object", + side_effect=all_gather_object, + ): + assert backend.initialize() is True + + assert set(backend._tensor_map) == {expert_name} + assert backend._receiver_session_ids == ["recv:7000"] + assert backend._session_debug_info["recv:7000"]["endpoint_idx"] == 0 + + def test_initialize_multi_sender_uses_explicit_sender_process_group(self): + process_group = object() + backend, _ = self._make_multi_sender_backend( + rank_index=0, + world_size=4, + rank_filter=lambda loc: True, + sender_ranks=(0, 2), + process_group=process_group, + ) + backend._receiver_session_ids = ["recv:7000"] + backend._last_prepare_returned_tensor_map = False + backend._initialize_single_sender = lambda: True # type: ignore[assignment] + all_gather_calls = [] + broadcast_groups = [] + + def all_gather_object(output, value, **kwargs): + assert len(output) == 2 + all_gather_calls.append((value, kwargs.get("group"))) + output[:] = [False, False] if len(all_gather_calls) == 1 else [True, True] + + def broadcast_payload(payload, src=0, **kwargs): + broadcast_groups.append(kwargs.get("group")) + payload[0] = ("reuse_cached", ["recv:7000"]) + + with patch("xorl.server.weight_sync.backends.p2p.dist.is_available", return_value=True): + with patch("xorl.server.weight_sync.backends.p2p.dist.is_initialized", return_value=True): + with patch( + "xorl.server.weight_sync.backends.p2p.dist.broadcast_object_list", + side_effect=broadcast_payload, + ): + with patch( + "xorl.server.weight_sync.backends.p2p.dist.all_gather_object", + side_effect=all_gather_object, + ): + assert backend.initialize() is True + + assert all(group is process_group for _, group in all_gather_calls) + assert broadcast_groups == [process_group] + + def test_adopt_prepared_state_skips_http(self): + backend, _ = self._make_multi_sender_backend(rank_index=2, world_size=4, rank_filter=lambda loc: True) + tensor_map = { + "model.layers.0.self_attn.q_proj.weight": [ + _hf_locator( + tp_rank=0, + full_shape=[128, 64], + slc=[[0, 64], [0, 64]], + ptr=0xAAAA0000, + nbytes=64 * 64 * 2, + session_id="recv:7000", + ), + ] + } + # Should not raise, should not POST anything. + ok = backend.adopt_prepared_state(tensor_map, ["recv:7000"]) + assert ok is True + assert backend._tensor_map == tensor_map + assert backend._receiver_session_ids == ["recv:7000"] + + def test_initialize_multi_sender_propagates_nonzero_rank_failure_to_rank0(self): + backend, _ = self._make_multi_sender_backend(rank_index=0, world_size=2, rank_filter=lambda loc: True) + + def initialize_rank0(): + backend._receiver_session_ids = ["recv:7000"] + backend._last_prepare_returned_tensor_map = False + return True + + all_gather_calls = [] + + def all_gather_object(output, value, **kwargs): + all_gather_calls.append(value) + if len(all_gather_calls) == 1: + output[:] = [False, False] + else: + output[:] = [True, False] + + backend._initialize_single_sender = initialize_rank0 # type: ignore[assignment] + + with patch("xorl.server.weight_sync.backends.p2p.dist.is_available", return_value=True): + with patch("xorl.server.weight_sync.backends.p2p.dist.is_initialized", return_value=True): + with patch("xorl.server.weight_sync.backends.p2p.dist.broadcast_object_list"): + with patch( + "xorl.server.weight_sync.backends.p2p.dist.all_gather_object", + side_effect=all_gather_object, + ): + assert backend.initialize() is False + + assert all_gather_calls == [False, True] + + def test_initialize_multi_sender_nonzero_rank_does_not_prewarm_engine_by_default(self): + backend, fake_engine = self._make_multi_sender_backend(rank_index=1, world_size=2, rank_filter=lambda loc: True) + backend._engine = None + tensor_map = { + "model.layers.0.self_attn.q_proj.weight": [ + _hf_locator( + tp_rank=0, + full_shape=[128, 64], + slc=[[0, 64], [0, 64]], + ptr=0xAAAA0000, + nbytes=64 * 64 * 2, + session_id="recv:7000", + ), + ] + } + events = [] + all_gather_calls = [] + + def make_local_engine(): + events.append("make_engine") + return fake_engine + + def broadcast_payload(payload, src=0, **kwargs): + events.append("broadcast") + payload[0] = ("tensor_map", tensor_map, ["recv:7000"]) + + def all_gather_object(output, value, **kwargs): + all_gather_calls.append(value) + output[:] = [False, False] if len(all_gather_calls) == 1 else [True, True] + + backend._make_local_engine = make_local_engine # type: ignore[assignment] + + with patch("xorl.server.weight_sync.backends.p2p.dist.is_available", return_value=True): + with patch("xorl.server.weight_sync.backends.p2p.dist.is_initialized", return_value=True): + with patch( + "xorl.server.weight_sync.backends.p2p.dist.broadcast_object_list", + side_effect=broadcast_payload, + ): + with patch( + "xorl.server.weight_sync.backends.p2p.dist.all_gather_object", + side_effect=all_gather_object, + ): + assert backend.initialize() is True + + assert events == ["broadcast", "make_engine"] + assert backend._engine is fake_engine + assert backend._tensor_map == tensor_map + assert backend._receiver_session_ids == ["recv:7000"] + + def test_initialize_multi_sender_nonzero_rank_prewarms_engine_before_broadcast(self, monkeypatch): + monkeypatch.setenv("XORL_P2P_PREINIT_NONZERO_ENGINES", "1") + backend, fake_engine = self._make_multi_sender_backend(rank_index=1, world_size=2, rank_filter=lambda loc: True) + backend._engine = None + tensor_map = { + "model.layers.0.self_attn.q_proj.weight": [ + _hf_locator( + tp_rank=0, + full_shape=[128, 64], + slc=[[0, 64], [0, 64]], + ptr=0xAAAA0000, + nbytes=64 * 64 * 2, + session_id="recv:7000", + ), + ] + } + events = [] + all_gather_calls = [] + + def make_local_engine(): + events.append("make_engine") + return fake_engine + + def broadcast_payload(payload, src=0, **kwargs): + events.append("broadcast") + payload[0] = ("tensor_map", tensor_map, ["recv:7000"]) + + def all_gather_object(output, value, **kwargs): + all_gather_calls.append(value) + output[:] = [False, False] if len(all_gather_calls) == 1 else [True, True] + + backend._make_local_engine = make_local_engine # type: ignore[assignment] + + with patch("xorl.server.weight_sync.backends.p2p.dist.is_available", return_value=True): + with patch("xorl.server.weight_sync.backends.p2p.dist.is_initialized", return_value=True): + with patch( + "xorl.server.weight_sync.backends.p2p.dist.broadcast_object_list", + side_effect=broadcast_payload, + ): + with patch( + "xorl.server.weight_sync.backends.p2p.dist.all_gather_object", + side_effect=all_gather_object, + ): + assert backend.initialize() is True + + assert events == ["make_engine", "broadcast"] + assert backend._engine is fake_engine + assert backend._tensor_map == tensor_map + assert backend._receiver_session_ids == ["recv:7000"] + + def test_rank_filter_routes_slices_to_owning_rank(self): + # Locators tagged with ep_rank; each sender ships only its own. + full = torch.arange(128 * 64, dtype=torch.bfloat16).reshape(128, 64) + locators = [ + { + **_hf_locator( + tp_rank=tp, + full_shape=[128, 64], + slc=[[tp * 64, (tp + 1) * 64], [0, 64]], + ptr=0xBEEF_0000 + tp * 0x10000, + nbytes=64 * 64 * 2, + session_id=f"recv-{tp}:7000", + ), + "hf_name": "model.layers.0.self_attn.q_proj.weight", + "ep_rank": tp, + "endpoint_idx": 0, + } + for tp in (0, 1) + ] + bucket = [("model.layers.0.self_attn.q_proj.weight", full)] + + rank0_engine_calls: List[Tuple[str, List[int], List[int], List[int]]] = [] + rank1_engine_calls: List[Tuple[str, List[int], List[int], List[int]]] = [] + + for rank_index, sink in ((0, rank0_engine_calls), (1, rank1_engine_calls)): + backend, engine = self._make_multi_sender_backend( + rank_index=rank_index, + world_size=2, + rank_filter=lambda loc, r=rank_index: loc.get("ep_rank") == r, + ) + backend._tensor_map = {locators[0]["hf_name"]: locators} + backend._receiver_session_ids = [locators[0]["session_id"], locators[1]["session_id"]] + backend.transfer_bucket(bucket, src_rank=rank_index) + backend.flush_pending_transfers() + sink.extend(engine.transfers) + + # Each rank issued exactly one transfer to its owning receiver. + assert len(rank0_engine_calls) == 1 + assert rank0_engine_calls[0][0] == "recv-0:7000" + assert len(rank1_engine_calls) == 1 + assert rank1_engine_calls[0][0] == "recv-1:7000" + + def test_transfer_bucket_accepts_nonzero_src_rank_when_direct_ep(self): + backend, _ = self._make_multi_sender_backend(rank_index=3, world_size=4, rank_filter=lambda loc: False) + backend._tensor_map = {} + # Should NOT raise even though src_rank != 0. + backend.transfer_bucket([], src_rank=3) + + def test_transfer_bucket_all_filtered_out_is_ok_for_direct_ep_rank(self): + backend, engine = self._make_multi_sender_backend(rank_index=3, world_size=4, rank_filter=lambda loc: False) + backend._tensor_map = { + "param": [ + _hf_locator( + tp_rank=0, + full_shape=[8, 8], + slc=[[0, 8], [0, 8]], + ptr=0x1000, + nbytes=8 * 8 * 2, + session_id="recv0:7000", + ) + ] + } + backend.transfer_bucket([("param", torch.zeros(8, 8, dtype=torch.bfloat16))], src_rank=3) + backend.flush_pending_transfers() + assert engine.transfers == [] + + +class TestP2PDestroy: + def test_complete_sync_does_not_complete_receiver_after_pending_transfer_failure(self): + backend, _ = _make_backend() + future: Future[None] = Future() + future.set_exception(RuntimeError("mooncake transfer failed")) + backend._cpu_pool_pending_futures[0] = future + + with patch("requests.post") as posted: + with pytest.raises(RuntimeError, match="mooncake transfer failed"): + backend.complete_sync() + + posted.assert_not_called() + assert backend._cpu_pool_pending_futures[0] is None + + def test_destroy_can_skip_receiver_completion_for_failed_sync_cleanup(self): + backend, engine = _make_backend() + backend._registered_source_ptrs = [0x1, 0x2] + + with patch("requests.post") as posted: + backend.destroy(complete_receiver=False) + + posted.assert_not_called() + assert engine.deregistered == [[0x1, 0x2]] + assert backend._engine is None + + def test_destroy_skip_completion_drains_failed_pending_transfers_without_raising(self): + backend, _ = _make_backend() + future: Future[None] = Future() + future.set_exception(RuntimeError("transfer failed during cleanup")) + backend._cpu_pool_pending_futures[0] = future + + with patch("requests.post") as posted: + backend.destroy(complete_receiver=False) + + posted.assert_not_called() + assert backend._cpu_pool_pending_futures == [None] * backend._n_pools + assert backend._engine is None + + def test_destroy_posts_complete_with_correct_payload(self): + backend, engine = _make_backend() + # Pretend transfer_bucket already ran and stashed flush/version state. + backend.config.backend_config["flush_cache"] = True + backend.config.backend_config["weight_version"] = "rev-7" + backend._registered_source_ptrs = [0x1, 0x2] + + with patch( + "requests.post", + return_value=_FakeResponse(200, {"success": True, "message": "ok"}), + ) as posted: + backend.destroy() + + # /complete_weights_update was called with the right body. + posted.assert_called_once() + url = posted.call_args.args[0] + body = posted.call_args.kwargs["json"] + assert url == "http://infer-0:5000/complete_weights_update" + assert body["transport"] == "p2p" + assert body["group_name"] == "weight_sync_group" + assert body["flush_cache"] is True + assert body["weight_version"] == "rev-7" + + # Source pointers were deregistered. + assert engine.deregistered == [[0x1, 0x2]] + + def test_destroy_raises_after_complete_failure_but_cleans_up(self): + backend, _ = _make_backend() + with patch("requests.post", return_value=_FakeResponse(500, {})): + with pytest.raises(RuntimeError, match="/complete_weights_update failed"): + backend.destroy() + assert backend._engine is None + + def test_complete_sync_defaults_flush_cache_false(self): + backend, _ = _make_backend() + with patch( + "requests.post", + return_value=_FakeResponse(200, {"success": True, "message": "ok"}), + ) as posted: + backend.complete_sync() + + body = posted.call_args.kwargs["json"] + assert body["flush_cache"] is False + + +class TestP2PEngineConstruction: + def test_resolve_local_hostname_prefers_pod_ip(self, monkeypatch): + monkeypatch.delenv("XORL_P2P_HOSTNAME", raising=False) + monkeypatch.setenv("POD_IP", "10.42.31.44") + monkeypatch.setenv("HOST_IP", "10.0.0.2") + + assert _resolve_local_hostname() == "10.42.31.44" + + def test_resolve_local_hostname_uses_socket_ip_before_fqdn(self, monkeypatch): + monkeypatch.delenv("XORL_P2P_HOSTNAME", raising=False) + monkeypatch.delenv("POD_IP", raising=False) + monkeypatch.delenv("HOST_IP", raising=False) + monkeypatch.delenv("HOSTNAME_IP", raising=False) + monkeypatch.setattr(socket, "gethostname", lambda: "trainer-pod") + monkeypatch.setattr(socket, "getfqdn", lambda: "trainer-pod.default.svc") + monkeypatch.setattr( + socket, + "gethostbyname", + lambda host: "10.42.31.44" if host == "trainer-pod" else "10.96.0.10", + ) + + assert _resolve_local_hostname() == "10.42.31.44" + + def test_make_local_engine_falls_back_without_sglang_wrapper(self, monkeypatch): + class FakeTransferEngine: + def initialize(self, hostname, protocol, transport, device): + self.initialized = (hostname, protocol, transport, device) + return 0 + + def get_rpc_port(self): + return 12345 + + def batch_register_memory(self, _ptrs, _lengths): + return 0 + + def batch_unregister_memory(self, _ptrs): + return 0 + + def batch_transfer_sync_write(self, _session_id, _buffers, _peer_buffer_addresses, _lengths): + return 0 + + mooncake_mod = types.ModuleType("mooncake") + mooncake_engine_mod = types.ModuleType("mooncake.engine") + mooncake_engine_mod.TransferEngine = FakeTransferEngine + monkeypatch.setitem(sys.modules, "mooncake", mooncake_mod) + monkeypatch.setitem(sys.modules, "mooncake.engine", mooncake_engine_mod) + monkeypatch.setitem(sys.modules, "sglang", None) + + cfg = TransportConfig( + endpoints=[EndpointConfig(host="infer-0", port=5000, world_size=1)], + group_name="weight_sync_group", + backend_config={"hostname": "trainer-0", "gpu_id": 0, "ib_device": "mlx5_0"}, + ) + backend = P2PTransportBackend(cfg) + engine = backend._make_local_engine() + + assert engine is not None + assert engine.get_session_id() == "trainer-0:12345" + assert engine.get_ib_device() == "mlx5_0" + assert engine.engine.initialized == ("trainer-0", "P2PHANDSHAKE", "rdma", "mlx5_0") + assert backend._engine is None diff --git a/tests/server/weight_sync/test_p2p_trainer_ib_device.py b/tests/server/weight_sync/test_p2p_trainer_ib_device.py new file mode 100644 index 00000000..ef24a760 --- /dev/null +++ b/tests/server/weight_sync/test_p2p_trainer_ib_device.py @@ -0,0 +1,140 @@ +import pytest + +from xorl.server.weight_sync import handler as handler_mod +from xorl.server.weight_sync.handler import WeightSyncHandler, _select_p2p_ib_device + + +def test_selects_global_rank_mapping_when_map_covers_world(monkeypatch): + monkeypatch.setenv("P2P_TRAINER_IB_DEVICES_PER_RANK", "mlx5_0;mlx5_1;mlx5_2;mlx5_3") + monkeypatch.setenv("LOCAL_RANK", "0") + + assert _select_p2p_ib_device(rank=2, world_size=4) == "mlx5_2" + + +def test_selects_local_rank_mapping_when_map_is_per_node(monkeypatch): + monkeypatch.setenv("P2P_TRAINER_IB_DEVICES_PER_RANK", "mlx5_0;mlx5_1;mlx5_2;mlx5_3") + monkeypatch.setenv("LOCAL_RANK", "1") + + assert _select_p2p_ib_device(rank=9, world_size=16) == "mlx5_1" + + +def test_falls_back_to_single_trainer_ib_device(monkeypatch): + monkeypatch.delenv("P2P_TRAINER_IB_DEVICES_PER_RANK", raising=False) + monkeypatch.delenv("P2P_TRAINER_GPU_TO_IB_DEVICE_MAP", raising=False) + monkeypatch.setenv("P2P_TRAINER_IB_DEVICE", "mlx5_6") + + assert _select_p2p_ib_device(rank=3, world_size=8) == "mlx5_6" + + +def test_selects_physical_gpu_mapping_from_visible_gpu_indices(monkeypatch): + monkeypatch.delenv("P2P_TRAINER_IB_DEVICES_PER_RANK", raising=False) + monkeypatch.delenv("P2P_TRAINER_IB_DEVICE", raising=False) + monkeypatch.setenv("P2P_TRAINER_GPU_TO_IB_DEVICE_MAP", "0=mlx5_2,1=mlx5_3,3=mlx5_5") + monkeypatch.setenv("P2P_TRAINER_VISIBLE_GPU_INDICES", "3,1,0") + + monkeypatch.setenv("LOCAL_RANK", "0") + assert _select_p2p_ib_device(rank=0, world_size=8) == "mlx5_5" + + monkeypatch.setenv("LOCAL_RANK", "1") + assert _select_p2p_ib_device(rank=1, world_size=8) == "mlx5_3" + + +def test_selects_physical_gpu_mapping_from_numeric_cuda_visible_devices(monkeypatch): + monkeypatch.delenv("P2P_TRAINER_IB_DEVICES_PER_RANK", raising=False) + monkeypatch.delenv("P2P_TRAINER_IB_DEVICE", raising=False) + monkeypatch.delenv("P2P_TRAINER_VISIBLE_GPU_INDICES", raising=False) + monkeypatch.delenv("SELECTED_GPU_INDICES", raising=False) + monkeypatch.setenv("P2P_TRAINER_GPU_TO_IB_DEVICE_MAP", "0:mlx5_2;6:mlx5_6") + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "6,0") + monkeypatch.setenv("LOCAL_RANK", "0") + + assert _select_p2p_ib_device(rank=0, world_size=8) == "mlx5_6" + + +def test_empty_mapping_entry_uses_autodiscovery(monkeypatch): + monkeypatch.setenv("P2P_TRAINER_IB_DEVICES_PER_RANK", "mlx5_0;;mlx5_2") + monkeypatch.setenv("LOCAL_RANK", "1") + monkeypatch.delenv("P2P_TRAINER_IB_DEVICE", raising=False) + + assert _select_p2p_ib_device(rank=1, world_size=3) is None + + +def test_sync_abort_marker_roundtrip(tmp_path): + class Trainer: + train_config = {"output_dir": str(tmp_path)} + + handler = WeightSyncHandler(rank=3, world_size=8, trainer=Trainer()) + path = handler._sync_abort_path("weight_sync_group", "iter-1") + + handler._mark_sync_abort(path, RuntimeError("transfer failed")) + + with pytest.raises(RuntimeError, match="rank=3: transfer failed"): + handler._raise_if_sync_aborted(path) + + handler._clear_sync_abort(path) + handler._raise_if_sync_aborted(path) + + +def test_p2p_rank_summary_includes_handler_transfer_totals(monkeypatch): + monkeypatch.setenv("LOCAL_RANK", "2") + + class Backend: + def stats_summary(self): + return {"total_bytes": 321, "num_buckets": 4, "main_thread_s": 0.2} + + handler = WeightSyncHandler(rank=6, world_size=8, trainer=None) + + summary = handler._build_p2p_rank_summary( + Backend(), + is_sender=True, + transfer_wall_s=0.5, + total_bytes=123, + num_parameters=7, + num_buckets=3, + ib_device="mlx5_2", + phase_s={"direct_ep_s": 0.4}, + ) + + assert summary["rank"] == 6 + assert summary["local_rank"] == 2 + assert summary["total_bytes"] == 123 + assert summary["num_parameters"] == 7 + assert summary["num_buckets"] == 3 + assert summary["has_transfers"] is True + + +def test_aggregate_p2p_transfer_totals_sums_all_rank_counters(): + total_bytes, num_parameters, num_buckets = WeightSyncHandler._aggregate_p2p_transfer_totals( + [ + {"rank": 0, "total_bytes": 100, "num_parameters": 4, "num_buckets": 1}, + {"rank": 1, "total_bytes": 300, "num_parameters": 6, "num_buckets": 2}, + {"rank": 2, "total_bytes": 0, "num_parameters": 0, "num_buckets": 0}, + ], + total_bytes=100, + num_parameters=4, + num_buckets=1, + ) + + assert total_bytes == 400 + assert num_parameters == 10 + assert num_buckets == 3 + + +def test_p2p_transfer_status_gather_reports_peer_failure(monkeypatch): + handler = WeightSyncHandler(rank=1, world_size=2, trainer=None) + + monkeypatch.setattr(handler_mod.dist, "is_available", lambda: True) + monkeypatch.setattr(handler_mod.dist, "is_initialized", lambda: True) + + def fake_all_gather_object(gathered, local_status): + gathered[0] = {"rank": 0, "ok": False, "error": "RuntimeError: transfer failed"} + gathered[1] = local_status + + monkeypatch.setattr(handler_mod.dist, "all_gather_object", fake_all_gather_object) + + statuses = handler._gather_p2p_transfer_statuses(None) + + assert statuses == [ + {"rank": 0, "ok": False, "error": "RuntimeError: transfer failed"}, + {"rank": 1, "ok": True}, + ] diff --git a/tests/test_arguments.py b/tests/test_arguments.py index 07d32bc6..6f68d29a 100644 --- a/tests/test_arguments.py +++ b/tests/test_arguments.py @@ -4,6 +4,7 @@ import torch import yaml +import xorl.arguments as arguments_module from xorl.arguments import Arguments, parse_args @@ -42,6 +43,41 @@ def test_parse_args_accepts_signsgd_from_yaml(tmp_path, monkeypatch): assert args.train.optimizer == "signsgd" assert args.train.optimizer_kwargs == {} + assert args.train.load_weights_mode == "grouped" + + +def test_parse_args_accepts_distsignsgd_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": "distsignsgd", + "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 == "distsignsgd" + assert args.train.optimizer_kwargs == {} def test_parse_args_wires_muon_kwargs_from_yaml(tmp_path, monkeypatch): @@ -64,6 +100,7 @@ def test_parse_args_wires_muon_kwargs_from_yaml(tmp_path, monkeypatch): "muon_ns_use_quack_kernels": False, "muon_gram_ns_num_restarts": 2, "muon_gram_ns_restart_iterations": [2], + "muon_grouped_gram_ns_fp32_byte_limit": 23, "muon_grad_dtype": "fp32", "muon_update_dtype": "fp32", "muon_force_momentum_path": True, @@ -87,7 +124,49 @@ def test_parse_args_wires_muon_kwargs_from_yaml(tmp_path, monkeypatch): 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_grouped_gram_ns_fp32_byte_limit"] == 23 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 + + +def test_parse_args_resolves_auto_checkpoint_before_skip_validation(tmp_path, monkeypatch): + resolved_checkpoint = str(tmp_path / "outputs" / "checkpoints" / "global_step_10") + 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"), + "load_weights_mode": "skip", + "load_checkpoint_path": "auto", + "repo_commit": "test-commit", + "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)]) + monkeypatch.setattr( + arguments_module, + "get_checkpoint_path", + lambda output_dir, is_local_rank0, ckpt_manager: resolved_checkpoint, + ) + + args = parse_args(Arguments) + + assert args.train.load_checkpoint_path == resolved_checkpoint diff --git a/tests/trainers/test_deepseek_v3_training_guards.py b/tests/trainers/test_deepseek_v3_training_guards.py new file mode 100644 index 00000000..8b0c47e5 --- /dev/null +++ b/tests/trainers/test_deepseek_v3_training_guards.py @@ -0,0 +1,105 @@ +from types import SimpleNamespace + +import pytest + +from xorl.models.auto import build_foundation_model +from xorl.models.transformers.deepseek_v3.configuration_deepseek_v3 import DeepseekV3Config +from xorl.models.transformers.deepseek_v3.modeling_deepseek_v3 import DeepseekV3ForCausalLM +from xorl.models.transformers.deepseek_v3.support import validate_deepseek_v3_tensor_parallelism +from xorl.trainers.model_builder import build_training_model + + +pytestmark = [pytest.mark.cpu] + + +def _tiny_config() -> DeepseekV3Config: + config = DeepseekV3Config( + vocab_size=32, + hidden_size=16, + intermediate_size=32, + moe_intermediate_size=8, + num_hidden_layers=1, + num_attention_heads=2, + num_key_value_heads=2, + n_shared_experts=1, + n_routed_experts=4, + routed_scaling_factor=1.0, + kv_lora_rank=4, + q_lora_rank=8, + qk_nope_head_dim=4, + qk_rope_head_dim=4, + v_head_dim=8, + n_group=2, + topk_group=1, + num_experts_per_tok=2, + first_k_dense_replace=0, + ) + config._attn_implementation = "eager" + config._moe_implementation = "eager" + config._activation_native = True + return config + + +def _patch_tiny_model(monkeypatch): + monkeypatch.setattr( + "xorl.trainers.model_builder.build_foundation_model", + lambda **kwargs: DeepseekV3ForCausalLM(_tiny_config()), + ) + monkeypatch.setattr("xorl.trainers.model_builder.helper.print_device_mem_info", lambda *args, **kwargs: None) + + +def test_build_foundation_model_rejects_train_router_for_deepseek(): + with pytest.raises(ValueError, match="train_router=True"): + build_foundation_model(_tiny_config(), train_router=True) + + +@pytest.mark.parametrize( + ("override_kwargs", "match"), + [ + ({"freeze_router": False}, "requires freeze_router=True"), + ({"freeze_router": True, "enable_qlora": True}, "enable_qlora=True"), + ({"freeze_router": True, "merge_qkv": False}, "merge_qkv=False"), + ], +) +def test_build_training_model_rejects_unsupported_deepseek_modes(monkeypatch, override_kwargs, match): + _patch_tiny_model(monkeypatch) + + kwargs = { + "config_path": "unused", + "weights_path": "unused", + "freeze_router": True, + "enable_mixed_precision": False, + "enable_gradient_checkpointing": False, + } + kwargs.update(override_kwargs) + + with pytest.raises(ValueError, match=match): + build_training_model(**kwargs) + + +def test_build_training_model_freezes_router_when_requested(monkeypatch): + _patch_tiny_model(monkeypatch) + monkeypatch.setattr("xorl.trainers.model_builder._parallelize", lambda model, **kwargs: model) + + result = build_training_model( + config_path="unused", + weights_path="unused", + freeze_router=True, + enable_mixed_precision=False, + enable_gradient_checkpointing=False, + ) + + router_params = [param for name, param in result.model.named_parameters() if ".gate.weight" in name] + assert router_params + assert all(param.requires_grad is False for param in router_params) + assert result.model.model.layers[0].self_attn.q_a_proj.weight.requires_grad is True + + +def test_validate_deepseek_v3_tensor_parallelism_rejects_tp(monkeypatch): + monkeypatch.setattr( + "xorl.models.transformers.deepseek_v3.support.get_parallel_state", + lambda: SimpleNamespace(tp_enabled=True), + ) + + with pytest.raises(ValueError, match="tensor parallelism is not supported yet"): + validate_deepseek_v3_tensor_parallelism(_tiny_config()) diff --git a/tests/trainers/test_lora_mixed_precision_dtype_policy.py b/tests/trainers/test_lora_mixed_precision_dtype_policy.py index 34154005..cd180952 100644 --- a/tests/trainers/test_lora_mixed_precision_dtype_policy.py +++ b/tests/trainers/test_lora_mixed_precision_dtype_policy.py @@ -32,11 +32,8 @@ def fake_parallelize(model, **kwargs): 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.build_foundation_model", fake_build_foundation_model) + monkeypatch.setattr("xorl.trainers.model_builder._parallelize", fake_parallelize) monkeypatch.setattr( "xorl.trainers.model_builder.helper.print_device_mem_info", lambda *args, **kwargs: None, diff --git a/tests/trainers/test_training_utils.py b/tests/trainers/test_training_utils.py new file mode 100644 index 00000000..31141af7 --- /dev/null +++ b/tests/trainers/test_training_utils.py @@ -0,0 +1,153 @@ +import math + +import pytest +import torch +import torch.nn as nn + +import xorl.trainers.training_utils as training_utils_module +from xorl.data.constants import IGNORE_INDEX +from xorl.trainers.training_utils import ( + clip_gradients, + count_active_microbatches, + get_distsign_grad_scale_factor, + get_effective_grad_clip_value, + sync_sp_gradients, +) + + +pytestmark = [pytest.mark.cpu] + + +class TinyModule(nn.Module): + def __init__(self): + super().__init__() + self.weight = nn.Parameter(torch.ones(2, dtype=torch.float32)) + + +def test_get_effective_grad_clip_value_preserves_regular_clipping(): + model = TinyModule() + model.weight.grad = torch.ones_like(model.weight) + + grad_norm = clip_gradients( + model, + get_effective_grad_clip_value(1.0, use_distsignsgd=False), + ) + + expected_scale = 1.0 / math.sqrt(2.0) + assert grad_norm == pytest.approx(math.sqrt(2.0)) + assert torch.allclose(model.weight.grad, torch.full_like(model.weight.grad, expected_scale)) + + +def test_get_effective_grad_clip_value_skips_clipping_for_distsignsgd(): + model = TinyModule() + model.weight.grad = torch.ones_like(model.weight) + + grad_norm = clip_gradients( + model, + get_effective_grad_clip_value(1.0, use_distsignsgd=True), + ) + + assert grad_norm == pytest.approx(math.sqrt(2.0)) + assert torch.equal(model.weight.grad, torch.ones_like(model.weight.grad)) + + +def test_get_distsign_grad_scale_factor_returns_mean_vote_scale(): + assert get_distsign_grad_scale_factor(8) == pytest.approx(0.125) + + +def test_get_distsign_grad_scale_factor_is_noop_without_active_voters(): + assert get_distsign_grad_scale_factor(0) == pytest.approx(1.0) + assert get_distsign_grad_scale_factor(-1) == pytest.approx(1.0) + + +def test_count_active_microbatches_batches_reduce_and_returns_voter_total(monkeypatch): + reduce_calls = [] + + def fake_all_reduce(tensor, op, group=None): + # Simulate a 4-rank DP group where rank counts (per-mb voter totals) are: + # mb0: 4 voters, mb1: 0 voters, mb2: 2 voters + # Local tensor here is the rank-0 contribution; SUM across the group + # would yield the per-mb voter counts above. + reduce_calls.append((tensor.clone(), op, group)) + tensor.copy_(torch.tensor([4, 0, 2], dtype=torch.int64)) + + monkeypatch.setattr(training_utils_module.dist, "all_reduce", fake_all_reduce) + + micro_batches = [ + {"labels": torch.tensor([1, 2, 3])}, + {"labels": torch.tensor([IGNORE_INDEX, IGNORE_INDEX])}, + {"labels": torch.tensor([4, IGNORE_INDEX, 5])}, + ] + + active_microbatches, active_voter_total = count_active_microbatches(micro_batches, group="dp") + + # Exactly one reduce, regardless of microbatch count. + assert len(reduce_calls) == 1 + _, op, group = reduce_calls[0] + assert op == torch.distributed.ReduceOp.SUM + assert group == "dp" + assert active_microbatches == 2 # mbs with at least one voter + assert active_voter_total == 6 # 4 + 0 + 2 + + +def test_count_active_microbatches_is_empty_input_safe(): + assert count_active_microbatches([]) == (0, 0) + + +def test_sync_sp_gradients_reduces_every_grad_by_default(monkeypatch): + reduced = [] + + class FakeParam: + def __init__(self, grad): + self.grad = grad + + class FakeModel: + def parameters(self): + return [ + FakeParam(torch.tensor([1.0, -2.0])), + FakeParam(torch.tensor([3.0, 4.0])), + ] + + def fake_all_reduce(tensor, op, group): + reduced.append((tensor.clone(), op, group)) + + monkeypatch.setattr(training_utils_module.dist, "all_reduce", fake_all_reduce) + + sync_sp_gradients(FakeModel(), sp_grad_sync_group="sp-group") + + assert [t.tolist() for t, _, _ in reduced] == [[1.0, -2.0], [3.0, 4.0]] + assert all(op == torch.distributed.ReduceOp.SUM for _, op, _ in reduced) + assert all(group == "sp-group" for _, _, group in reduced) + + +def test_sync_sp_gradients_skips_dtensor_grads_when_requested(monkeypatch): + reduced = [] + + class FakeDTensor: + def __init__(self, local_tensor): + self._local_tensor = local_tensor + + class FakeParam: + def __init__(self, grad): + self.grad = grad + + class FakeModel: + def parameters(self): + return [ + FakeParam(FakeDTensor(torch.tensor([1.0, -2.0]))), + FakeParam(torch.tensor([3.0, 4.0])), + ] + + def fake_all_reduce(tensor, op, group): + reduced.append((tensor.clone() if isinstance(tensor, torch.Tensor) else tensor, op, group)) + + monkeypatch.setattr(training_utils_module, "DTensor", FakeDTensor) + monkeypatch.setattr(training_utils_module.dist, "all_reduce", fake_all_reduce) + + sync_sp_gradients(FakeModel(), sp_grad_sync_group="sp-group", skip_dtensor_grads=True) + + assert len(reduced) == 1 + tensor, op, group = reduced[0] + assert torch.equal(tensor, torch.tensor([3.0, 4.0])) + assert op == torch.distributed.ReduceOp.SUM + assert group == "sp-group" diff --git a/tests/utils/test_distillation_teacher_cache.py b/tests/utils/test_distillation_teacher_cache.py new file mode 100644 index 00000000..346e0d1b --- /dev/null +++ b/tests/utils/test_distillation_teacher_cache.py @@ -0,0 +1,137 @@ +import pytest +import torch + +from tests._helpers.opd import save_tensor_file +from xorl.distillation import ( + TeacherActivationCache, + TeacherHeadManager, + TeacherHeadStore, + load_lm_head_weight, + prepare_lm_head_teacher_store, +) + + +def test_load_lm_head_weight_from_safetensors_file(tmp_path): + weight = torch.randn(11, 7) + path = tmp_path / "teacher_head.safetensors" + save_tensor_file(path, "lm_head.weight", weight) + + loaded = load_lm_head_weight(str(path)) + + torch.testing.assert_close(loaded, weight) + + +def test_teacher_head_manager_keeps_one_device_head(tmp_path): + weight0 = torch.randn(5, 3) + weight1 = torch.randn(5, 3) + path0 = tmp_path / "teacher0.safetensors" + path1 = tmp_path / "teacher1.safetensors" + save_tensor_file(path0, "lm_head.weight", weight0) + save_tensor_file(path1, "lm_head.weight", weight1) + + manager = TeacherHeadManager({"0": str(path0), "1": str(path1)}) + + loaded0 = manager.get(0, device="cpu") + loaded1 = manager.get(1, device="cpu") + + torch.testing.assert_close(loaded0, weight0) + torch.testing.assert_close(loaded1, weight1) + assert manager._device_teacher_id == "1" + + +def test_teacher_head_manager_reloads_for_dtype_change(tmp_path): + weight = torch.randn(5, 3) + path = tmp_path / "teacher.safetensors" + save_tensor_file(path, "lm_head.weight", weight) + + manager = TeacherHeadManager({"0": str(path)}) + + loaded_fp32 = manager.get(0, device="cpu", dtype=torch.float32) + loaded_bf16 = manager.get(0, device="cpu", dtype=torch.bfloat16) + + assert loaded_fp32.dtype == torch.float32 + assert loaded_bf16.dtype == torch.bfloat16 + torch.testing.assert_close(loaded_bf16.float(), weight, rtol=1e-2, atol=1e-2) + + +def test_teacher_store_round_trips_lm_head_shards(tmp_path): + weight = torch.randn(11, 7) + model_dir = tmp_path / "teacher_model" + model_dir.mkdir() + save_tensor_file(model_dir / "model.safetensors", "lm_head.weight", weight) + + store_dir = tmp_path / "teacher_store" + manifest = prepare_lm_head_teacher_store(model_dir, store_dir, teacher_id=3, shard_rows=4) + store = TeacherHeadStore(manifest) + + loaded = store.load_lm_head(3) + via_loader = load_lm_head_weight(str(store_dir), teacher_id=3) + + torch.testing.assert_close(loaded, weight) + torch.testing.assert_close(via_loader, weight) + assert [shard.rows for shard in store.head_spec(3).shards] == [4, 4, 3] + + +def test_teacher_head_manager_prefetches_teacher_store(tmp_path): + weight = torch.randn(5, 3) + model_dir = tmp_path / "teacher_model" + model_dir.mkdir() + save_tensor_file(model_dir / "model.safetensors", "lm_head.weight", weight) + store_dir = tmp_path / "teacher_store" + prepare_lm_head_teacher_store(model_dir, store_dir, teacher_id=0, shard_rows=2) + + manager = TeacherHeadManager({"0": str(store_dir)}) + manager.prefetch(0) + loaded = manager.get(0, device="cpu") + + torch.testing.assert_close(loaded, weight) + + +def test_teacher_activation_cache_gathers_indices(tmp_path): + hidden = torch.randn(6, 4) + path = tmp_path / "hidden_states.safetensors" + save_tensor_file(path, "hidden_states", hidden) + + cache = TeacherActivationCache({"3": str(path)}) + indices = torch.tensor([[0, 2, 5], [1, 4, 3]]) + + gathered = cache.get(3, indices, device="cpu") + + torch.testing.assert_close(gathered, hidden[indices]) + + +def test_teacher_activation_cache_prefetches(tmp_path): + hidden = torch.randn(6, 4) + path = tmp_path / "hidden_states.safetensors" + save_tensor_file(path, "hidden_states", hidden) + + cache = TeacherActivationCache({"3": str(path)}) + cache.prefetch(3) + gathered = cache.get(3, torch.tensor([0, 5]), device="cpu") + + torch.testing.assert_close(gathered, hidden[[0, 5]]) + + +def test_teacher_activation_cache_rejects_negative_indices(tmp_path): + """Negative indices used to be silently clamped to 0, masking producer bugs.""" + hidden = torch.randn(6, 4) + path = tmp_path / "hidden_states.safetensors" + save_tensor_file(path, "hidden_states", hidden) + + cache = TeacherActivationCache({"3": str(path)}) + bad_indices = torch.tensor([[0, 2, -1], [1, 4, 3]]) + + with pytest.raises(IndexError, match="negative"): + cache.get(3, bad_indices, device="cpu") + + +def test_teacher_activation_cache_rejects_out_of_range_indices(tmp_path): + hidden = torch.randn(6, 4) + path = tmp_path / "hidden_states.safetensors" + save_tensor_file(path, "hidden_states", hidden) + + cache = TeacherActivationCache({"3": str(path)}) + bad_indices = torch.tensor([0, 2, 6]) + + with pytest.raises(IndexError, match="cache only has 6 rows"): + cache.get(3, bad_indices, device="cpu")