From e1fbf4ebdf306109dcf0134fe4b9cfe8913c7c03 Mon Sep 17 00:00:00 2001 From: Conner Manuel <57027354+connermanuel@users.noreply.github.com> Date: Mon, 20 Apr 2026 09:10:26 -0700 Subject: [PATCH 01/49] Lazy import rotary pos emb in strategy --- README.md | 8 ++++---- docs/src/content/docs/getting-started/installation.md | 4 ++-- docs/src/content/docs/getting-started/quickstart.md | 6 +++--- .../docs/server-training/client-sdk/loss-functions.md | 4 ++-- .../server-training/examples/password-memorization.md | 4 ++-- .../docs/server-training/examples/sft-no-robots.md | 2 +- src/xorl/distributed/sequence_parallel/strategy.py | 4 +++- 7 files changed, 17 insertions(+), 15 deletions(-) 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/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/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/src/xorl/distributed/sequence_parallel/strategy.py b/src/xorl/distributed/sequence_parallel/strategy.py index 96342f50..1beaa791 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 @@ -309,6 +308,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 From fc310575accc5c9bd90e8c1f0accb1d27d991aaa Mon Sep 17 00:00:00 2001 From: Qingyang Wu Date: Mon, 20 Apr 2026 14:40:57 -0700 Subject: [PATCH 02/49] Fix dummy dataset packing bins race on non-zero ranks The dummy-dataset fast path in prepare_datasets called PackingDataset on all ranks concurrently, violating the rank-0-first + barrier contract that PR established for _load_or_compute_bins. Non-zero ranks would race rank 0's bin cache write and fail after 10 retries. Gate the dummy path with the same pattern used for real datasets: rank 0 builds the dataset (computing and caching bins), barrier, then non-zero ranks load from cache. --- src/xorl/data/prepare/prepare_datasets.py | 23 +++++++++++++++++++---- 1 file changed, 19 insertions(+), 4 deletions(-) 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 From 029b1ace7b37bf48215d2f995d85defe3f718da9 Mon Sep 17 00:00:00 2001 From: Qingyang Wu Date: Tue, 21 Apr 2026 12:39:43 -0700 Subject: [PATCH 03/49] Delete unused direct_train CLI entry point src/xorl/cli/direct_train.py was an unreferenced secondary training entry point. Keeping it alive means every change to training internals risks drifting between it and the primary xorl.cli.train path. --- src/xorl/cli/direct_train.py | 844 ----------------------------------- 1 file changed, 844 deletions(-) delete mode 100644 src/xorl/cli/direct_train.py 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() From 7f77597c38473cf87d8e3ce07a9ba7ef2dfc5cc1 Mon Sep 17 00:00:00 2001 From: Qingyang Wu Date: Wed, 22 Apr 2026 10:52:05 -0700 Subject: [PATCH 04/49] Raise NotImplementedError instead of falling back to HuggingFace * Raise NotImplementedError instead of silently falling back to HuggingFace When an arch was not in the xorl registry, the loader used to hand the model to AutoModel/AutoModelForCausalLM with only a warning. If an xorl modeling module failed to import, that produced a silent fallback in which no xorl code paths ran, confusing downstream debugging (see). The registry now records per-module import errors, and the loader raises NotImplementedError (surfacing those errors) rather than falling back. * Apply ruff-format --- src/xorl/models/loader.py | 25 ++++++++++++++----------- src/xorl/models/registry.py | 4 +++- 2 files changed, 17 insertions(+), 12 deletions(-) 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/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 From 8d64504ddacf62e61ad9f2bf3e4a8f620a5d2125 Mon Sep 17 00:00:00 2001 From: qywu Date: Wed, 22 Apr 2026 13:44:38 -0700 Subject: [PATCH 05/49] MoE backend improvements: configurable activation, fused GEMM, deterministic accumulation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * [Feat] MoE backend improvements: configurable activation, fused GEMM, deterministic accumulation Three improvements to all MoE backends (native, triton, quack) that improve correctness, performance, and generality: 1. **Configurable activation function** β€” thread `act_fn` from MoEExperts through all backends and EP paths. Previously hardcoded to F.silu(). Uses compile-friendly `use_gelu_tanh` flag for torch.compile paths. Backward uses generic autograd grad for non-silu activations. 2. **Fused gate_up GEMM** β€” single `grouped_mm(x, gate_up_proj)` then chunk, instead of two separate gate/up GEMMs. ~9% faster, eliminates bf16 rounding difference between fused and split matmuls. 3. **Post-GEMM routing weights** β€” apply routing weights (expert_scores) after the down projection GEMM, not before. Mathematically equivalent (scalar-per-row commutes with matmul) but produces different bf16 rounding. Critical for models with high-magnitude norm weights. 4. **Deterministic accumulation** (triton/quack local path) β€” replace moe_gather's index_add_ with reshape(tokens, top_k, hidden).sum(). 5x faster (no GPU atomics), deterministic across runs. Changes per file: - experts.py: pass act_fn + gate_up_weight to backend functions - backend/__init__.py: EP compute wrappers (kwargsβ†’positional for apply()) - backend/native.py: fused GEMM path, post-GEMM routing, reshape-sum - backend/triton.py, quack.py: pass **kwargs through to inner functions - ops/moe/triton.py: all 4 autograd Functions updated (forward+backward) - ops/moe/quack.py: all 5 autograd Functions updated (forward+backward) Verified: - Qwen3 MoE (silu): max diff 3e-5 across backends (no regression) - Gemma4 MoE (gelu_pytorch_tanh): bit-exact across all 3 backends - EP (2 GPU): triton==quack, both activations, forward+backward OK * style: apply ruff/ruff-format fixes Address lint CI failures in test files and merged MoE backend files. * fix: align XORL_DEBUG_EP quack path with post-GEMM routing; drop gemma4 smoke tests Quack's _forward_debug applied expert_scores before the down GEMM while the normal path applies it after. Different bf16 rounding meant debug mode changed numerics. Move the multiplication to match the normal path. Also remove tests/test_gemma4_* smoke scripts β€” they call dist.init_process_group at import time and break `pytest tests/` collection, and the gemma4 module they depend on isn't in this repo. * test: drop remaining EP smoke scripts that break pytest collection test_qwen3_ep_quick.py and test_ep_debug.py call dist.init_process_group at import time and fail pytest collection. Remove them along with their _import_fix helper (only referenced by the deleted files). * refactor: drop dead moe_act variants Remove classes/functions that are defined but never selected by the backend registry or called anywhere: TritonEPGroupGemmMoeAct, TritonMoeExpertsFunctionMoeAct, triton_moe_forward_moe_act, QuackEPGroupGemmMoeAct, QuackMoeExpertsFunctionMoeAct, quack_moe_forward_moe_act, and native's _run_experts_moe_act / _gate_up_swiglu. * refactor(moe): thread act_kind string instead of act_fn callable Replace fragile `_is_gelu_tanh(act_fn)` name-sniffing with an explicit `act_kind: str` threaded from MoEExperts through the triton/native/quack backends and their ops-layer autograd Functions. - Add `normalize_act_kind` (maps e.g. "gelu_pytorch_tanh" -> "gelu_tanh") and `check_act_kind_supported` in ops/moe/triton.py as the shared source. - Each backend declares SUPPORTED_ACT_KINDS and validates at entry β€” unsupported activations now raise ValueError instead of silently falling back to SiLU. - MoEExperts stores self.act_kind (normalized) alongside self.act_fn (still used by the eager and LoRA paths). * refactor(moe): rename act_kind -> hidden_act for config consistency Matches the upstream HF `config.hidden_act` vocabulary users already know. The value is still normalized ('gelu_pytorch_tanh' -> 'gelu_tanh') via `normalize_hidden_act` so backends see a canonical set. * refactor(moe): drop dead split-path branch in TritonMoeExpertsFunction MoEExperts always stores weights as a single fused gate_up_proj and passes it to the local path, so the `else` branch (two separate gate/up GEMMs) was unreachable. Assert gate_up_weight is set and remove the dead branch in both forward and backward. Note: QuackMoeExpertsFunction (and native_ep_compute) still use the split-path pattern β€” those are separate refactors since they would need a new fused GEMM implementation, not just branch removal. * refactor(moe): make hidden_act explicit in backend shims Both backend/triton.py and backend/quack.py previously relied on `**kwargs` as a pass-through channel for hidden_act. The docstrings said kwargs were "ignored" which was wrong β€” hidden_act was consumed downstream. Make hidden_act (and gate_up_weight) explicit parameters. Same change in the ops-layer `triton_moe_forward` / `quack_moe_forward` entry points. * refactor(moe): use fused gate_up GEMM in quack and native EP paths Makes the PR's "fused gate+up GEMM" claim honest for all three backends, and standardizes the kwarg name on `gate_up_proj` (previously `gate_up_weight` for the local path, split `gate_proj`/`up_proj` for the quack/native EP paths). - QuackMoeExpertsFunction: single GEMM on gate_up_proj, fused dgrad/wgrad in backward (mirrors TritonMoeExpertsFunction). - QuackEPGroupGemm: takes gate_up_proj directly, fused dgrad/wgrad. - native_ep_compute: takes gate_up_proj and routes through _run_experts_grouped_mm's fused branch. - _quack_ep_fused / _native_ep_fused shims: stop slicing +.contiguous() copying the fused weight. - Rename gate_up_weight -> gate_up_proj across callers and autograd Functions for consistency with MoEExperts.gate_up_proj. QuackTPMoeExpertsFunction (TP path) still uses split GEMMs and would need a separate refactor. * test(ep): update test_adapter_source to match fused gate_up_proj signature The source-string regression test was asserting on the old pre-fuse signature (_QuackEPGroupGemm.apply(permute_tokens, cumsum, gate_proj, up_proj, down_proj, expert_scores)). After making quack EP use a single fused gate_up_proj GEMM, the call site now uses (permute_tokens, cumsum, gate_up_proj, down_proj, intermediate_size, expert_scores, hidden_act). Update the expected substrings to match. Same for the native EP wrapper. --------- --- .../models/layers/moe/backend/__init__.py | 28 +- src/xorl/models/layers/moe/backend/native.py | 142 ++++--- src/xorl/models/layers/moe/backend/quack.py | 10 +- src/xorl/models/layers/moe/backend/triton.py | 12 +- src/xorl/models/layers/moe/experts.py | 8 + src/xorl/ops/moe/quack.py | 391 ++++++++++-------- src/xorl/ops/moe/triton.py | 270 +++++++----- tests/ops/test_ep_adapter_wrappers.py | 19 +- 8 files changed, 534 insertions(+), 346 deletions(-) diff --git a/src/xorl/models/layers/moe/backend/__init__.py b/src/xorl/models/layers/moe/backend/__init__.py index 5b8f0796..34d929d7 100644 --- a/src/xorl/models/layers/moe/backend/__init__.py +++ b/src/xorl/models/layers/moe/backend/__init__.py @@ -52,7 +52,14 @@ 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" + ): + 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 @@ -60,10 +67,12 @@ 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" + ): + 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: @@ -73,10 +82,11 @@ def _quack_ep_fused(permute_tokens, cumsum, gate_up_proj, down_proj, intermediat try: from .native import native_ep_compute as _native_ep_compute - def _native_ep_fused(permute_tokens, cumsum, gate_up_proj, down_proj, intermediate_size, 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" + ): + del intermediate_size + return _native_ep_compute(permute_tokens, cumsum, gate_up_proj, down_proj, expert_scores, hidden_act=hidden_act) EP_EXPERT_COMPUTE["native"] = _native_ep_fused except ImportError: diff --git a/src/xorl/models/layers/moe/backend/native.py b/src/xorl/models/layers/moe/backend/native.py index 118fb395..6ea7a361 100644 --- a/src/xorl/models/layers/moe/backend/native.py +++ b/src/xorl/models/layers/moe/backend/native.py @@ -128,55 +128,80 @@ 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, ) -> 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:: - - gate_proj: [num_experts, hidden_dim, intermediate_size] - up_proj: [num_experts, hidden_dim, intermediate_size] - down_proj: [num_experts, intermediate_size, hidden_dim] + 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. """ 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 + 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 hidden_act == "gelu_tanh": + gate_out = F.gelu(gate_raw, approximate="tanh") + else: + gate_out = F.silu(gate_raw) + + # GLU: act(gate) * up h = gate_out * up_out - if expert_scores is not None: - h = h * expert_scores.to(h.dtype).unsqueeze(-1) # down: h @ down_proj -> (tokens, hidden) + # expert_scores applied AFTER down GEMM (not before) for bf16 consistency out = torch._grouped_mm( h, down_proj.to(compute_dtype), offs=offsets, ).to(x.dtype) + 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 # --------------------------------------------------------------------------- @@ -199,6 +224,7 @@ def _native_expert_forward_impl( down_proj: torch.Tensor, num_experts: int, compute_fn, + gate_up_proj: torch.Tensor | None = None, ) -> torch.Tensor: """Shared token-sort / pad / scatter logic for native expert forward. @@ -245,25 +271,28 @@ 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) β€” NO expert_scores inside GEMM expert_out_padded = compute_fn( gate_proj, up_proj, down_proj, sorted_hidden_padded, padded_counts, - expert_scores_padded, + None, # expert_scores applied after, not inside GEMM + gate_up_proj=gate_up_proj, ) - # 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 +302,18 @@ def _native_expert_forward_impl( # --------------------------------------------------------------------------- +SUPPORTED_HIDDEN_ACTS = frozenset({"silu", "gelu_tanh"}) + + +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) + return _run_experts_compiled + + def native_expert_forward( hidden_states: torch.Tensor, routing_weights: torch.Tensor, @@ -281,9 +322,14 @@ 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``.""" + 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 +338,8 @@ def native_expert_forward( up_proj, down_proj, num_experts, - _run_experts_compiled, + _select_compiled_fn(hidden_act), + gate_up_proj=gate_up_proj, ) @@ -496,34 +543,24 @@ 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", ) -> 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]``. """ + 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 +570,12 @@ 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 + 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 ) - # 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..94af2a93 100644 --- a/src/xorl/models/layers/moe/experts.py +++ b/src/xorl/models/layers/moe/experts.py @@ -71,6 +71,10 @@ 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) # EP dispatch strategy: "alltoall" (default) or "deepep" (NVLink-optimized) self.ep_dispatch: str = "alltoall" @@ -140,6 +144,8 @@ def forward( up_proj, self.down_proj, num_experts=self.num_experts, + hidden_act=self.hidden_act, + gate_up_proj=self.gate_up_proj, ) @torch.compiler.disable @@ -235,6 +241,7 @@ def _ep_forward( self.down_proj, self.intermediate_size, expert_scores, + hidden_act=self.hidden_act, ) # Step 3: Combine expert outputs back to original ranks @@ -268,6 +275,7 @@ 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, ) ev[3].record() diff --git a/src/xorl/ops/moe/quack.py b/src/xorl/ops/moe/quack.py index 9237c14e..7ed1de10 100644 --- a/src/xorl/ops/moe/quack.py +++ b/src/xorl/ops/moe/quack.py @@ -6,6 +6,7 @@ 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.triton import _moe_gate_activation, _moe_gate_activation_backward, check_hidden_act_supported def _debug_ep_enabled() -> bool: @@ -25,52 +26,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 +91,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 +146,184 @@ 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 - ) - 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 + 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, ) - - # wgrad FC1 - grad_gate_proj = None - if gate_proj.requires_grad: - grad_gate_proj = torch.empty_like(gate_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_gate_output, - c=grad_gate_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_gate_output - grad_up_proj = None - if up_proj.requires_grad: - grad_up_proj = torch.empty_like(up_proj) - quack_group_gemm_same_mn( - a=scatter_output, - b=grad_up_output, - c=grad_up_proj, - cumsum_K=cumsum_t, - max_K=max_M, - transpose_a=True, - transpose_b=False, - cu_seqlens_k=cu_seqlens_m, - ) - del grad_up_output, scatter_output + 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 + grad_expert_scores = None + if ctx.has_expert_scores: + 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 +331,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 +421,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 +474,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 +511,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 +557,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 +579,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 +592,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..de440c8c 100644 --- a/src/xorl/ops/moe/triton.py +++ b/src/xorl/ops/moe/triton.py @@ -1,6 +1,6 @@ import torch +import torch.nn.functional as F -from xorl.ops.fused_silu_and_mul import silu_and_mul, silu_and_mul_backward from xorl.utils.import_utils import is_fused_moe_available @@ -8,26 +8,70 @@ 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, ) +# Canonical activation kinds understood by MoE ops. Upstream `hidden_act` +# strings (e.g. "gelu_pytorch_tanh") are normalized to one of these. +SUPPORTED_HIDDEN_ACTS: frozenset[str] = frozenset({"silu", "gelu_tanh"}) + + +def normalize_hidden_act(hidden_act: str | None) -> str: + """Normalize a HF-style ``hidden_act`` string to an 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" + 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 _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 +81,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 +110,29 @@ 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 + grad_expert_scores = None + if ctx.has_expert_scores: + 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 +144,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 +196,7 @@ def backward(ctx, grad_output): grad_down_proj, None, # intermediate_size grad_expert_scores, + None, # hidden_act ) @@ -145,6 +208,8 @@ class TritonMoeExpertsFunction(torch.autograd.Function): to free dead tensors immediately. """ + SUPPORTED_HIDDEN_ACTS = frozenset({"silu", "gelu_tanh"}) + @staticmethod def forward( ctx, @@ -155,63 +220,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 +275,7 @@ def forward( cumsum_t, gate_output, up_output, - scattered_gate_weight, + gate_up_proj, ) return output @@ -239,21 +292,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 +319,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 +330,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 +339,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 +400,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 +416,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 +430,6 @@ def triton_moe_forward( gate_proj, up_proj, down_proj, - gate_up_weight, + gate_up_proj, + hidden_act, ) diff --git a/tests/ops/test_ep_adapter_wrappers.py b/tests/ops/test_ep_adapter_wrappers.py index 94e90053..b9dd5f04 100644 --- a/tests/ops/test_ep_adapter_wrappers.py +++ b/tests/ops/test_ep_adapter_wrappers.py @@ -159,7 +159,7 @@ 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) @@ -173,9 +173,16 @@ def test_adapter_source_forwards_expert_scores(): """ source = _BACKEND_INIT_PATH.read_text() - assert "_QuackEPGroupGemm.apply(permute_tokens, cumsum, gate_proj, up_proj, down_proj, expert_scores)" in source, ( - "_quack_ep_fused does not forward expert_scores to _QuackEPGroupGemm.apply()" - ) - assert "_native_ep_compute(permute_tokens, cumsum, gate_proj, up_proj, down_proj, expert_scores)" in source, ( - "_native_ep_fused does not forward expert_scores to _native_ep_compute()" + assert "expert_scores" in source and "_QuackEPGroupGemm.apply(" in source, ( + "_QuackEPGroupGemm.apply() call missing from backend/__init__.py" ) + # Quack EP shim must forward expert_scores (and hidden_act) through. + assert ( + "_QuackEPGroupGemm.apply(\n permute_tokens, cumsum, gate_up_proj, down_proj, intermediate_size, expert_scores, hidden_act\n )" + in source + ), "_quack_ep_fused does not forward expert_scores to _QuackEPGroupGemm.apply()" + # Native EP shim must forward expert_scores (and hidden_act) through. + assert ( + "_native_ep_compute(permute_tokens, cumsum, gate_up_proj, down_proj, expert_scores, hidden_act=hidden_act)" + in source + ), "_native_ep_fused does not forward expert_scores to _native_ep_compute()" From 6586ab7306be35875bdaac08019d2980a665afbe Mon Sep 17 00:00:00 2001 From: Conner Manuel <57027354+connermanuel@users.noreply.github.com> Date: Wed, 22 Apr 2026 17:48:59 -0700 Subject: [PATCH 06/49] Fix: set default grad checkpointing correctly * set default grad chkpting * linter --- src/xorl/models/base.py | 8 +++++--- src/xorl/models/module_utils.py | 8 +++++--- 2 files changed, 10 insertions(+), 6 deletions(-) diff --git a/src/xorl/models/base.py b/src/xorl/models/base.py index 2d5b76a5..ce3ebf73 100644 --- a/src/xorl/models/base.py +++ b/src/xorl/models/base.py @@ -91,8 +91,10 @@ 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. + # Default to "recompute_full_layer" so downstream consumers can rely on + # `_gradient_checkpointing_method` always holding a valid mode string. + grad_ckpt_method = gradient_checkpointing_kwargs.pop("gradient_checkpointing_method", "recompute_full_layer") grad_ckpt_func = partial(torch.utils.checkpoint.checkpoint, **gradient_checkpointing_kwargs) @@ -102,7 +104,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 != "recompute_full_layer": 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/module_utils.py b/src/xorl/models/module_utils.py index 6ac45c14..2f625cdc 100644 --- a/src/xorl/models/module_utils.py +++ b/src/xorl/models/module_utils.py @@ -1591,12 +1591,13 @@ class GradientCheckpointingLayer(nn.Module): """ gradient_checkpointing = False + _gradient_checkpointing_method = "recompute_full_layer" 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 == "recompute_full_layer" ): return self._gradient_checkpointing_func(partial(super().__call__, **kwargs), *args) return super().__call__(*args, **kwargs) @@ -1621,6 +1622,7 @@ class MoEGradientCheckpointingLayer(nn.Module): """ gradient_checkpointing = False + _gradient_checkpointing_method = "recompute_full_layer" def _pre_mlp_forward(self, hidden_states, **kwargs): """Layernorm β†’ attention β†’ layernorm. Override per model. @@ -1659,8 +1661,8 @@ def _moe_forward(self, hidden_states, output_router_logits=False, **kwargs): _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 != "recompute_full_layer" ) _is_moe = isinstance(self.mlp, MoEBlock) From 9d2a2e531cccb542fc99cc20837d33c7f165831f Mon Sep 17 00:00:00 2001 From: Conner Manuel <57027354+connermanuel@users.noreply.github.com> Date: Fri, 24 Apr 2026 09:16:19 -0700 Subject: [PATCH 07/49] Add grad chkpt tests and defaults --- src/xorl/models/base.py | 13 ++- src/xorl/models/module_utils.py | 16 +++- .../qwen3_5_moe/modeling_qwen3_5_moe.py | 2 +- .../qwen3_moe/modeling_qwen3_moe.py | 6 +- tests/models/test_gradient_checkpointing.py | 92 +++++++++++++++++++ 5 files changed, 117 insertions(+), 12 deletions(-) create mode 100644 tests/models/test_gradient_checkpointing.py diff --git a/src/xorl/models/base.py b/src/xorl/models/base.py index ce3ebf73..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 @@ -92,9 +94,12 @@ def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs=None): gradient_checkpointing_kwargs = {"use_reentrant": False} # Pop selective checkpoint config before passing to torch checkpoint. - # Default to "recompute_full_layer" so downstream consumers can rely on - # `_gradient_checkpointing_method` always holding a valid mode string. - grad_ckpt_method = gradient_checkpointing_kwargs.pop("gradient_checkpointing_method", "recompute_full_layer") + # 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) @@ -104,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 != "recompute_full_layer": + 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/module_utils.py b/src/xorl/models/module_utils.py index 2f625cdc..ca934df8 100644 --- a/src/xorl/models/module_utils.py +++ b/src/xorl/models/module_utils.py @@ -1567,6 +1567,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,13 +1599,13 @@ class GradientCheckpointingLayer(nn.Module): """ gradient_checkpointing = False - _gradient_checkpointing_method = "recompute_full_layer" + _gradient_checkpointing_method: GradientCheckpointingMethod = DEFAULT_GRADIENT_CHECKPOINTING_METHOD def __call__(self, *args, **kwargs): if ( self.gradient_checkpointing and self.training - and self._gradient_checkpointing_method == "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) @@ -1622,7 +1630,7 @@ class MoEGradientCheckpointingLayer(nn.Module): """ gradient_checkpointing = False - _gradient_checkpointing_method = "recompute_full_layer" + _gradient_checkpointing_method: GradientCheckpointingMethod = DEFAULT_GRADIENT_CHECKPOINTING_METHOD def _pre_mlp_forward(self, hidden_states, **kwargs): """Layernorm β†’ attention β†’ layernorm. Override per model. @@ -1662,7 +1670,7 @@ def _moe_forward(self, hidden_states, output_router_logits=False, **kwargs): _selective = ( self.training and self.gradient_checkpointing - and self._gradient_checkpointing_method != "recompute_full_layer" + and self._gradient_checkpointing_method != DEFAULT_GRADIENT_CHECKPOINTING_METHOD ) _is_moe = isinstance(self.mlp, MoEBlock) 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..4ce0e1d9 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 @@ -522,7 +522,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_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/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 From 16015183e2ec344df0bc823810c76805b1d1e4e6 Mon Sep 17 00:00:00 2001 From: Qingyang Wu Date: Tue, 28 Apr 2026 09:07:16 -0700 Subject: [PATCH 08/49] Add sglang_shared_outer LoRA export format * Add sglang_shared_outer LoRA export format Teach save_lora_checkpoint to emit hybrid-shared MoE LoRA directly in SGLang's shared_outer layout (stacked 3D tensors, w1/w2/w3 slots, out-first dim order), and extend the PEFT load path to accept the same keys so roundtrip works. Avoids the external _lora_convert/convert.py re-pack step when targeting SGLang. * Apply ruff-format to lora utils * Wire LoRA export format through server checkpoint save * Export MoE LoRA in PEFT orientation by default * Fix hybrid-shared MoE LoRA export for SGLang * Mirror moe_hybrid_shared_lora and force collective gather for MoE LoRA - Write `moe_hybrid_shared_lora: True` into adapter_config.json when `lora_export_format="sglang_shared_outer"`. SGLang's lora_manager classifies adapters by this flag, and shared_outer is hybrid_shared on disk; without it SGLang mis-classifies as per_expert and rejects loading under `--lora-moe-format hybrid_shared`. - Skip the fast adapter-manager LoRA save path for any MoE LoRA. The fast path reads rank-local params without an EP all_gather and would export only num_local_experts (e.g. 32) instead of num_experts (e.g. 128), silently truncating the saved adapter. --------- Co-authored-by: Ashwinee Panda --- src/xorl/lora/utils.py | 78 +++++++++-- src/xorl/server/runner/checkpoint/manager.py | 15 +- src/xorl/server/server_arguments.py | 8 ++ .../runner/test_lora_checkpoint_roundtrip.py | 132 ++++++++++++++++++ 4 files changed, 221 insertions(+), 12 deletions(-) diff --git a/src/xorl/lora/utils.py b/src/xorl/lora/utils.py index bf601e08..fdc04fbd 100644 --- a/src/xorl/lora/utils.py +++ b/src/xorl/lora/utils.py @@ -37,6 +37,13 @@ _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") def _get_submodule(model: nn.Module, target: str) -> Tuple[nn.Module, str]: @@ -435,6 +442,17 @@ 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) + continue + match = _MOE_PEFT_LORA_PATTERN.match(key) if match is not None: prefix, expert_token, proj_name, lora_type = match.groups() @@ -478,7 +496,8 @@ 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", ) -> str: """ Save LoRA weights in PEFT-compatible format. @@ -499,13 +518,27 @@ 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``). 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 @@ -582,6 +615,21 @@ def _unmerge_moe_lora_weights(name: str, stacked_tensor: torch.Tensor) -> Dict[s 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 = stacked_tensor.transpose(-2, -1).contiguous().to(torch.bfloat16) + return peft_key, out_tensor + # Convert keys to PEFT format: base_model.model.{converted_key} peft_state_dict = {} detected_modules = set() @@ -590,18 +638,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(".") @@ -649,8 +699,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: diff --git a/src/xorl/server/runner/checkpoint/manager.py b/src/xorl/server/runner/checkpoint/manager.py index f411a31d..ee05a664 100644 --- a/src/xorl/server/runner/checkpoint/manager.py +++ b/src/xorl/server/runner/checkpoint/manager.py @@ -148,17 +148,27 @@ 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: it reads rank-local params without an + # all_gather across EP ranks, so it would export only num_local_experts + # (e.g. 32) instead of num_experts (e.g. 128) for per-expert MoE tensors. + is_moe_lora = bool(self.lora_config.get("moe_hybrid_shared_lora", False)) or any( + m in (self.lora_config.get("lora_target_modules") or []) for m in ("gate_proj", "up_proj", "down_proj") + ) + 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() + lora_export_format = self.lora_config.get("lora_export_format", "peft") save_lora_checkpoint( model=self.model, save_path=save_path, @@ -168,6 +178,7 @@ def _save_lora_weights(self, save_path: str, model_id: str) -> None: 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, ) # Cleanup diff --git a/src/xorl/server/server_arguments.py b/src/xorl/server/server_arguments.py index c2f452ab..0ee1f4fb 100644 --- a/src/xorl/server/server_arguments.py +++ b/src/xorl/server/server_arguments.py @@ -448,6 +448,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 # ======================================================================== @@ -601,6 +608,7 @@ def to_config_dict(self) -> Dict[str, Any]: "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, diff --git a/tests/server/runner/test_lora_checkpoint_roundtrip.py b/tests/server/runner/test_lora_checkpoint_roundtrip.py index 13c1f9fd..8cb7b550 100644 --- a/tests/server/runner/test_lora_checkpoint_roundtrip.py +++ b/tests/server/runner/test_lora_checkpoint_roundtrip.py @@ -1,9 +1,11 @@ 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 @@ -80,6 +82,59 @@ def _actual_lora_state(module: nn.Module) -> dict[str, torch.Tensor]: return {name: param.detach().cpu().clone() for name, param in _iter_lora_parameters(module)} +def test_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 +160,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) From 22cc1327160456266ce4b7700c6c1b7de69abb64 Mon Sep 17 00:00:00 2001 From: Ashwinee Panda Date: Tue, 28 Apr 2026 17:12:25 -0300 Subject: [PATCH 09/49] Allow qkv_proj in LoRA attention target_modules whitelist --- src/xorl/lora/utils.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/xorl/lora/utils.py b/src/xorl/lora/utils.py index fdc04fbd..a6d55c19 100644 --- a/src/xorl/lora/utils.py +++ b/src/xorl/lora/utils.py @@ -907,7 +907,9 @@ def inject_lora_into_model_with_moe( 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"]] + attention_modules = [ + m for m in target_modules if m in ["q_proj", "k_proj", "v_proj", "qkv_proj", "o_proj", "lm_head"] + ] expert_modules = [m for m in target_modules if m in ["gate_proj", "up_proj", "down_proj"]] # Step 1: Inject LoRA into standard nn.Linear layers From ba1754c7f997bae5f34ad2b123bfc8a19837dfa8 Mon Sep 17 00:00:00 2001 From: Ashwinee Panda Date: Wed, 29 Apr 2026 13:11:00 -0700 Subject: [PATCH 10/49] Forward lr_min and lr_decay_ratio to linear LR schedule * Forward lr_min and lr_decay_ratio to linear LR schedule build_lr_scheduler dropped lr_min and lr_decay_ratio when lr_decay_style was "linear", so the linear schedule always decayed to ~0 over the full training horizon, regardless of the configured floor or decay ratio. Pass both through, and update get_linear_schedule_with_warmup to honor lr_decay_ratio (matching the cosine path) and to interpolate down to min_lr instead of 0. * Validate LR scheduler args and clean up cosine - Reject lr <= 0 and lr_warmup_ratio outside [0, 1] in build_lr_scheduler; the lambdas divide by init_lr and a non-positive lr would silently NaN. - Drop the dead max(0, factor) clamp and now-redundant assert in the cosine lambda; with min_lr_ratio >= 0 the factor is already non-negative. Hoist lr_decay_steps out of the lambda to match the linear path. - Add tests/optim/test_lr_scheduler.py covering constant/linear/cosine shapes, lr_min floor, lr_decay_ratio behavior, and the new validation. --- src/xorl/optim/lr_scheduler.py | 35 ++++--- tests/optim/test_lr_scheduler.py | 153 +++++++++++++++++++++++++++++++ 2 files changed, 175 insertions(+), 13 deletions(-) create mode 100644 tests/optim/test_lr_scheduler.py 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/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") From 978ff4edd80bea695b2d345dcdaaf577ecf8ce6e Mon Sep 17 00:00:00 2001 From: Qingyang Wu Date: Wed, 29 Apr 2026 14:05:49 -0700 Subject: [PATCH 11/49] Fix failing tests and add GPU CI workflow MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Fix stale and dead tests - test_ep_plan_and_shard_tensor: patch get_parallel_state where parallel_plan imports it, not at source - test_ep_group_gemm_propagates_routing_score_gradients[quack]: pass fused gate_up_proj (backend API was unified in); add moe_add_gather to the kernel stub so the test works in isolation - TestBackendGKN.test_backends_match_reference_and_agree: pass fused gate_up_proj (Triton MoE now asserts it is provided) - test_lora_mixed_precision_keeps_base_bf16_and_skips_generic_upcast: patch build_foundation_model / _parallelize on model_builder where they are looked up - Delete tests/ops/test_moe_act.py: _moe_act flag is never read in src; all 36 tests compared identical computations, so they were vacuous - Delete tests/qlora/test_quantize_error_reduction.py: empty placeholder - Delete TestHubAndRemoteDetection in test_shared.py: always skipped because s3fs is not installed * Drop 4-GPU and 8-GPU vocab parallel CE variants test_vocab_parallel_ce_2gpu already exercises the same code path. Removing the larger variants drops the non-e2e suite's max GPU requirement from 8 to 2, letting it fit on a 2-GPU runner with far more concurrency (24 vs 6 slots org-wide). * Add GPU test workflow on self-hosted-h100-2gpu-cu131 Runs pytest -m "not e2e" inside nvcr.io/nvidia/cuda:12.9.1-cudnn-runtime-ubuntu24.04 on the existing org-scoped ARC scale set on turbox. 2 GPUs cover the full non-e2e suite after the vocab_parallel_ce trim. Image choice: CUDA 12.9.1 runtime matches the cu129 torch wheel pinned in pyproject.toml; cudnn-runtime variant includes cuDNN, NCCL, and nvidia-smi; Ubuntu 24.04 ships Python 3.12 which satisfies requires-python. ~3 GB vs ~25 GB for the pytorch container, so cold starts are minutes faster. git is installed at job start because actions/checkout@v4 requires it and the base image is a clean CUDA runtime. * Apply ruff-format to test_moe_gkn_format.py * Install build-essential for Triton JIT The nvcr cudnn-runtime image ships without a C compiler. Triton needs gcc (or cc) to compile the host-side kernel launcher on first use; without it, 108 GPU tests fail with 'Failed to find C compiler'. * Remove pr-test-cpu.yml in favor of GPU workflow The GPU workflow runs pytest -m 'not e2e' which already covers every CPU-marked test. Running them twice (once on dp-cp arc-runner-set and once on turbox H100 pod) wastes H100 time without adding signal. Also drop the arc-runner-set label from actionlint.yaml since no workflow references it anymore. * Restore CPU test workflow and address GPU workflow review CPU lane (arc-runner-set) is restored verbatim so cheap, fast feedback is preserved for PRs that don't need GPU signal β€” the H100 lane was ~4x slower and contended. GPU workflow review changes: - Switch base image to nvcr.io/nvidia/cuda:12.9.1-cudnn-devel-ubuntu24.04 so gcc/g++/make are present and build-essential apt install is dropped. - Use actions/checkout@v4 and astral-sh/setup-uv@v8 instead of pinned SHAs. - Enable setup-uv caching at /root/.cache/uv. - uv sync --frozen --group test (uv.lock is checked in and verified in sync). - Drop the nvidia-smi debug step. Re-add arc-runner-set to actionlint.yaml since the CPU workflow uses it. * Pin setup-uv to v8.0.0 (no moving v8 tag exists) * Drop uv sync --frozen (uv.lock is gitignored) * Skip cpu-marked tests in GPU lane to avoid duplicate runs CPU lane already runs pytest -m "cpu". With the GPU lane on "not e2e" they double-cover cpu-marked tests on contended H100 capacity. Narrow the GPU lane to "not e2e and not cpu" so each test runs in exactly one lane. * Disable setup-uv GH Actions cache; rely on arc-runner-cache PVC The runner pod already mounts /root/.cache from a PVC (arc-runner-cache), so uv's default cache directory persists across runs without any GH Actions cache round-trip. Keeping enable-cache: true triggered a post-job 'uv cache prune --ci --force' that consistently fails on the ARC k8s container with 'Executing the custom container implementation failed', marking otherwise-green runs as cancelled. --- tests/data/prepare/test_shared.py | 36 -- .../test_ep_lora_weight_slicing.py | 2 +- tests/distributed/test_vocab_parallel_ce.py | 16 - tests/ops/test_ep_routing_scores.py | 32 +- tests/ops/test_moe_act.py | 398 ------------------ tests/ops/test_moe_gkn_format.py | 5 +- tests/qlora/test_quantize_error_reduction.py | 6 - .../test_lora_mixed_precision_dtype_policy.py | 7 +- 8 files changed, 18 insertions(+), 484 deletions(-) delete mode 100644 tests/ops/test_moe_act.py delete mode 100644 tests/qlora/test_quantize_error_reduction.py 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_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_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/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_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..c0673f96 100644 --- a/tests/ops/test_moe_gkn_format.py +++ b/tests/ops/test_moe_gkn_format.py @@ -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/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/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, From 78526c012779da5efe15afbadfddd21fc31fb42b Mon Sep 17 00:00:00 2001 From: Qingyang Wu Date: Wed, 29 Apr 2026 14:28:32 -0700 Subject: [PATCH 12/49] ci: add conventional commit check on PR title Enforces Conventional Commits on PR titles (which become squash-merge commit messages). Comments the detected version bump on the PR, or fails the check with guidance if the title is invalid. Prerequisite for migrating the release workflow to semantic-release, where commit types drive the version bump. --- CONTRIBUTING.md | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) 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) From 21e54b4c04c74294dfac5df249d09375ca8456ff Mon Sep 17 00:00:00 2001 From: Ashwinee Panda Date: Wed, 29 Apr 2026 15:57:31 -0700 Subject: [PATCH 13/49] Add softmax auxiliary (Z-)loss on LM-head logits --- src/xorl/arguments.py | 24 +++ src/xorl/ops/loss/__init__.py | 8 +- src/xorl/ops/loss/causallm_loss.py | 74 +++++++-- src/xorl/ops/loss/compiled_cross_entropy.py | 51 ++++++ src/xorl/server/server_arguments.py | 4 +- src/xorl/trainers/trainer.py | 12 +- tests/ops/loss/test_causallm_z_loss.py | 172 ++++++++++++++++++++ 7 files changed, 331 insertions(+), 14 deletions(-) create mode 100644 tests/ops/loss/test_causallm_z_loss.py diff --git a/src/xorl/arguments.py b/src/xorl/arguments.py index 27c8a10d..7e1f1a37 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 @@ -664,6 +665,29 @@ def optimizer_kwargs(self) -> Dict[str, Any]: 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."}, diff --git a/src/xorl/ops/loss/__init__.py b/src/xorl/ops/loss/__init__.py index 18a8b620..ec7eed53 100644 --- a/src/xorl/ops/loss/__init__.py +++ b/src/xorl/ops/loss/__init__.py @@ -8,7 +8,7 @@ - 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 @@ -18,6 +18,11 @@ 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 # --------------------------------------------------------------------------- @@ -43,6 +48,7 @@ def register_loss_function(name: str, fn: Callable) -> None: __all__ = [ + "CrossEntropyMode", "LossOutput", "LOSS_REGISTRY", "get_loss_function", diff --git a/src/xorl/ops/loss/causallm_loss.py b/src/xorl/ops/loss/causallm_loss.py index 29bed248..cd9d643f 100644 --- a/src/xorl/ops/loss/causallm_loss.py +++ b/src/xorl/ops/loss/causallm_loss.py @@ -3,7 +3,10 @@ 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.vocab_parallel_cross_entropy import vocab_parallel_cross_entropy @@ -19,6 +22,7 @@ def causallm_loss_function( tp_group=None, use_compile: bool = False, lm_head_fp32: bool = False, + z_loss_coef: float = 0.0, ) -> "LossOutput": """ Compute causal language modeling loss. @@ -38,9 +42,19 @@ 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). + 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 @@ -52,6 +66,11 @@ def causallm_loss_function( # 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 @@ -73,40 +92,73 @@ def causallm_loss_function( ) return LossOutput(loss=loss) + z_loss_enabled = z_loss_coef > 0.0 + valid_count = valid_mask.sum().clamp(min=1) + 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 = per_token_ce.sum() / valid_count + if z_loss_enabled: + z_loss = per_token_lse_sq.sum() / valid_count + 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 = per_token_ce.sum() / valid_count + if z_loss_enabled: + z_loss = per_token_lse_sq.sum() / valid_count + 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..2e74868e 100644 --- a/src/xorl/ops/loss/compiled_cross_entropy.py +++ b/src/xorl/ops/loss/compiled_cross_entropy.py @@ -14,6 +14,9 @@ # 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] = {} + # Check if auto_chunker is available _AUTO_CHUNKER_AVAILABLE = None @@ -74,6 +77,31 @@ 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 _get_compiled_ce_fn(num_chunks: int, reduction: str = "none") -> Callable: """ Get or create a compiled cross-entropy function. @@ -107,3 +135,26 @@ 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] diff --git a/src/xorl/server/server_arguments.py b/src/xorl/server/server_arguments.py index 0ee1f4fb..7c0bdbba 100644 --- a/src/xorl/server/server_arguments.py +++ b/src/xorl/server/server_arguments.py @@ -13,6 +13,8 @@ import yaml +from xorl.ops.loss import CrossEntropyMode + @dataclass class ServerArguments: @@ -224,7 +226,7 @@ class ServerArguments: 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)" diff --git a/src/xorl/trainers/trainer.py b/src/xorl/trainers/trainer.py index dc58df5c..e0083036 100644 --- a/src/xorl/trainers/trainer.py +++ b/src/xorl/trainers/trainer.py @@ -231,6 +231,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: @@ -921,7 +931,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 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, + ) From 417048a761ae3a155e9dee829b5a664868ed9c2a Mon Sep 17 00:00:00 2001 From: Ashwinee Panda Date: Thu, 30 Apr 2026 07:58:00 -0700 Subject: [PATCH 14/49] Add DeepSeek V4 distributed Muon paths MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Add DeepSeek V4 distributed Muon paths Implements two opt-in features from DeepSeek V4 Β§3.5.1, both default off: 1. Full-gradient Muon NS (muon_distributed_mode=full_gradient) - Replaces shard-local NS on FSDP2/EP DTensor params with NS on the all-gathered full matrix, recovering the exact Muon update direction. - Momentum/Nesterov stay on the local shard (linear in grad, commutes with sharding) so the optimizer-state buffer stays at local-shard size. - LR adjustment uses the global matrix shape (fixes a bug where adjust_lr_fn saw the local-shard shape under shard_local mode). - Knapsack matrix-to-rank assignment and bounded ZeRO width are not in this PR; every rank in the param's mesh runs NS redundantly. 2. BF16 stochastic-rounded a2a + FP32 local sum (moe_grad_reduce_mode =bf16_a2a_fp32_sum) - Custom FSDPModule.set_custom_reduce_scatter on EP-experts modules: stochastic-round FP32β†’BF16, dist.all_to_all_single, sum locally in FP32. Halves comm volume vs FP32 reduce-scatter while preserving FP32 accumulation precision. - Stochastic rounding (xorl.optim.stochastic_round) operates on the FP32 bit pattern; unbiased in expectation. Tests: - CPU: stochastic round dtype, unbiasedness, bracket bounds, determinism. - Distributed (2 GPU): full-gradient Muon matches a single-rank oracle; shard-local provably differs on multi-rank. - Distributed (4 GPU): BF16 a2a reduce-scatter matches FP32 reduce-scatter within BF16 ulp bounds, with unbiased mean over many trials. - Distributed (2 GPU): end-to-end FSDP2 hook smoke test. * Address self-review: BF16 hook preconditions, plan typing, 3D EP test - Assert mp_policy.reduce_dtype=fp32 and gradient_divide_factor=1.0 on every expert FSDPModule before installing BF16StochasticAllToAllReduceScatter. Both are correctness preconditions: a non-fp32 reduce_dtype would surface a runtime error during the first backward (far from configuration), and a factor != 1.0 means FSDP applies predivide before the hook and skips the postdivide because the reduce is custom β€” silently under-weighting grads. Document the same preconditions on the hook's docstring. - Type _MuonUpdatePlan.placements as Optional[Tuple[Placement,...]] and device_mesh as Optional[DeviceMesh] via TYPE_CHECKING imports, replacing the loose Optional[object] / Optional[Tuple] placeholders. - Add a 3D-EP-experts (Shard(1) on the dp mesh) variant to the existing full_gradient distributed test, exercising the deferred-reshape branch in _muon_step that the previous tests only covered for 2D Shard(0). Both full_gradient (matches single-rank oracle) and shard_local (provably differs) layouts are checked. * Apply ruff-format to test files post-merge --- src/xorl/arguments.py | 24 ++ src/xorl/distributed/fsdp2/__init__.py | 1 + src/xorl/distributed/fsdp2/bf16_a2a_reduce.py | 93 ++++++++ src/xorl/distributed/torch_parallelize.py | 62 +++++- src/xorl/optim/muon.py | 104 ++++++++- src/xorl/optim/optimizer.py | 4 +- src/xorl/optim/stochastic_round.py | 43 ++++ src/xorl/server/runner/model_runner.py | 2 + src/xorl/server/server_arguments.py | 24 ++ src/xorl/trainers/model_builder.py | 2 + src/xorl/trainers/trainer.py | 1 + tests/distributed/test_bf16_a2a_fsdp_hook.py | 128 +++++++++++ tests/distributed/test_bf16_a2a_reduce.py | 131 +++++++++++ tests/distributed/test_muon_full_gradient.py | 208 ++++++++++++++++++ tests/optim/test_stochastic_round.py | 72 ++++++ 15 files changed, 891 insertions(+), 8 deletions(-) create mode 100644 src/xorl/distributed/fsdp2/bf16_a2a_reduce.py create mode 100644 src/xorl/optim/stochastic_round.py create mode 100644 tests/distributed/test_bf16_a2a_fsdp_hook.py create mode 100644 tests/distributed/test_bf16_a2a_reduce.py create mode 100644 tests/distributed/test_muon_full_gradient.py create mode 100644 tests/optim/test_stochastic_round.py diff --git a/src/xorl/arguments.py b/src/xorl/arguments.py index 7e1f1a37..35a1f3ee 100644 --- a/src/xorl/arguments.py +++ b/src/xorl/arguments.py @@ -630,6 +630,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]: @@ -659,6 +670,8 @@ 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( @@ -867,6 +880,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."}, 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..fb3fe979 --- /dev/null +++ b/src/xorl/distributed/fsdp2/bf16_a2a_reduce.py @@ -0,0 +1,93 @@ +"""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 + + +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") + if op != ReduceOp.SUM and op != ReduceOp.AVG: + raise NotImplementedError(f"BF16StochasticAllToAllReduceScatter requires SUM or AVG op, got {op}") + 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 == ReduceOp.AVG: + summed.div_(world_size) + output_tensor.copy_(summed) + return None diff --git a/src/xorl/distributed/torch_parallelize.py b/src/xorl/distributed/torch_parallelize.py index e92dbca1..2c63acbd 100644 --- a/src/xorl/distributed/torch_parallelize.py +++ b/src/xorl/distributed/torch_parallelize.py @@ -12,7 +12,7 @@ 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, @@ -93,6 +93,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 +216,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 +235,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,6 +333,60 @@ 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" @@ -438,6 +495,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": """ @@ -593,6 +651,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": @@ -731,6 +790,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/optim/muon.py b/src/xorl/optim/muon.py index 1f84e1f2..95d13761 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,14 +24,20 @@ 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 .gram_newton_schulz import GramNewtonSchulzOrthogonalizer, expand_ns_coefficients, find_best_restarts @@ -41,12 +47,39 @@ GROUPED_GRAM_NS_FP32_BYTE_LIMIT = 2 * 1024**3 +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 class _MuonUpdatePlan: param: torch.Tensor adjusted_lr: float orig_shape: Optional[torch.Size] pieces: list[Optional[torch.Tensor]] + placements: Optional[Tuple["Placement", ...]] = None + device_mesh: Optional["DeviceMesh"] = None @dataclass @@ -108,6 +141,17 @@ class Muon(TorchMuon): (``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__( @@ -130,6 +174,7 @@ def __init__( gram_newton_schulz_num_restarts: int = 1, gram_newton_schulz_restart_iterations: Optional[Iterable[int]] = None, adamw_state_dtype: Optional[torch.dtype] = None, + distributed_mode: str = "shard_local", ): if ns_algorithm not in {"standard_newton_schulz", "gram_newton_schulz"}: raise ValueError( @@ -140,12 +185,17 @@ 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'." + ) 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 @@ -201,8 +251,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"] @@ -224,6 +281,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 +298,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 +350,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 +381,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) @@ -335,6 +422,11 @@ 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) diff --git a/src/xorl/optim/optimizer.py b/src/xorl/optim/optimizer.py index 84ddf988..7a1aceae 100644 --- a/src/xorl/optim/optimizer.py +++ b/src/xorl/optim/optimizer.py @@ -174,6 +174,7 @@ 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, + distributed_mode=kwargs.get("muon_distributed_mode", "shard_local"), ) return Muon, ctor_kwargs else: @@ -356,7 +357,8 @@ def build_optimizer( "muon_ns_algorithm": "standard_newton_schulz", "muon_ns_use_quack_kernels": True, "muon_gram_ns_num_restarts": 1, "muon_gram_ns_restart_iterations": None, "muon_momentum_dtype": None, - "muon_grad_dtype": None, "muon_update_dtype": None, "muon_force_momentum_path": False} + "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 """ # EP-aware routing: for FSDP2+EP, split params into EP and non-EP groups and build two optimizers. 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/server/runner/model_runner.py b/src/xorl/server/runner/model_runner.py index dec1ef9c..7725947c 100644 --- a/src/xorl/server/runner/model_runner.py +++ b/src/xorl/server/runner/model_runner.py @@ -503,6 +503,7 @@ def _initialize_model(self): enable_forward_prefetch=self.train_config.get("enable_forward_prefetch", True), load_weights_mode=self.train_config.get("load_weights_mode", "broadcast"), 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), @@ -577,6 +578,7 @@ def _initialize_optimizer(self): "muon_grad_dtype", "muon_update_dtype", "muon_force_momentum_path", + "muon_distributed_mode", ) if k in self.train_config } diff --git a/src/xorl/server/server_arguments.py b/src/xorl/server/server_arguments.py index 7c0bdbba..5ce6118c 100644 --- a/src/xorl/server/server_arguments.py +++ b/src/xorl/server/server_arguments.py @@ -333,6 +333,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 # ======================================================================== @@ -588,6 +610,8 @@ def to_config_dict(self) -> Dict[str, Any]: "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, "load_checkpoint_path": self.load_checkpoint_path, "ckpt_manager": self.ckpt_manager, "enable_self_test": self.enable_self_test, diff --git a/src/xorl/trainers/model_builder.py b/src/xorl/trainers/model_builder.py index 39644534..497fad5b 100644 --- a/src/xorl/trainers/model_builder.py +++ b/src/xorl/trainers/model_builder.py @@ -132,6 +132,7 @@ def build_training_model( enable_forward_prefetch: bool = True, load_weights_mode: str = "broadcast", reshard_after_forward: Optional[bool] = None, + moe_grad_reduce_mode: str = "reduce_scatter", pp_schedule: Optional[str] = None, # --- Training flags --- freeze_router: bool = False, @@ -264,6 +265,7 @@ def build_training_model( load_weights_mode=load_weights_mode, pp_schedule=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, diff --git a/src/xorl/trainers/trainer.py b/src/xorl/trainers/trainer.py index e0083036..3fd79905 100644 --- a/src/xorl/trainers/trainer.py +++ b/src/xorl/trainers/trainer.py @@ -533,6 +533,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, 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..f3057b3e --- /dev/null +++ b/tests/distributed/test_bf16_a2a_reduce.py @@ -0,0 +1,131 @@ +"""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 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 _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_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/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) From 4ec734aabab14b99ec3ed3cfcd8f664a0d96f5b5 Mon Sep 17 00:00:00 2001 From: Qingyang Wu Date: Thu, 30 Apr 2026 12:45:11 -0700 Subject: [PATCH 15/49] fix(lora): merge in fp32 and cast the sum once MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * LoRA: merge in fp32 and cast the sum once Current ``LoraLinear.merge_weights`` and ``MoEExpertsLoRA.merge_weights`` round the LoRA delta per element to the weight dtype before adding: W.add_(delta.to(W.dtype)) This loses ~1 weight-dtype ULP per element of Ξ”. On bf16 weights with MoE top-k routing downstream, tiny per-element drift can flip expert selection at the router (argmax top-k is discontinuous), which cascades through the stack. Empirically, on Qwen3-30B-A3B with random lora_B ~N(0, 0.005) the naive variant degrades K3 by ~20x vs the fp32-sum-then-cast variant: std=0.005 naive K3=1.9e-1 fp32-cast-once K3=1.0e-2 (18x better) The change is a one-liner: upcast W to fp32, add the fp32 Ξ”, cast the sum back once. W.data.copy_((W.to(fp32) + delta).to(W.dtype)) Same memory as before (result lands in W.dtype). Strictly β‰₯ precision of the naive variant β€” bit-exact in fp32, more faithful in bf16/fp16. On realistic trained LoRAs (password-memorization adapters, std << 0.005) both variants produce identical greedy outputs; this change widens the "safe" margin and is the right default. New tests (tests/models/test_lora_merge_fp32_cast_once.py): - zero-LoRA merge is bit-exact - merged weight matches the fp32-sum-then-cast reference bit-for-bit - fp32-cast-once error ≀ naive error (vs true fp32 merged value) - same invariants on MoEExpertsLoRA (fused gate_up + down) All existing LoRA tests still pass (54 total). * style: apply ruff fixes to lora merge test --------- Co-authored-by: Ashwinee Panda --- src/xorl/lora/modules/linear.py | 12 +- src/xorl/models/layers/moe/lora.py | 11 +- .../models/test_lora_merge_fp32_cast_once.py | 140 ++++++++++++++++++ 3 files changed, 159 insertions(+), 4 deletions(-) create mode 100644 tests/models/test_lora_merge_fp32_cast_once.py diff --git a/src/xorl/lora/modules/linear.py b/src/xorl/lora/modules/linear.py index c6c4772f..b7510468 100644 --- a/src/xorl/lora/modules/linear.py +++ b/src/xorl/lora/modules/linear.py @@ -195,10 +195,18 @@ def merge_weights(self) -> None: Used both for periodic merge during training (merge_lora_interval) and one-shot merge for inference. + + Precision note: we upcast ``weight`` to float32, add the fp32 delta, + then cast the sum back β€” a *single* quantization at the end. This is + strictly more faithful than the naive ``weight += delta.to(W.dtype)`` + (which rounds Ξ” per element before adding) and matters in precision- + sensitive settings like MoE top-k routing on bf16 weights. Same memory + as the naive variant β€” both land in ``self.weight.dtype``. """ with torch.no_grad(): - delta_weight = self.get_delta_weight() - self.weight.add_(delta_weight.to(self.weight.dtype)) + delta_weight = self.get_delta_weight() # fp32 + 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: diff --git a/src/xorl/models/layers/moe/lora.py b/src/xorl/models/layers/moe/lora.py index 837ad78b..1550267a 100644 --- a/src/xorl/models/layers/moe/lora.py +++ b/src/xorl/models/layers/moe/lora.py @@ -180,14 +180,21 @@ def merge_weights(self) -> None: After merging: weight = weight + delta_weight for each active projection. Resets LoRA parameters after merge. + + Precision note: ``base`` is upcast to float32, the fp32 delta is added, + and the sum is cast back once. This is strictly more faithful than + rounding Ξ” per element before adding, and keeps unmerged-forward close + to merged-forward on MoE models where top-k routing amplifies any + per-element delta quantization error. """ with torch.no_grad(): for proj_name in ("gate_proj", "up_proj", "down_proj"): if proj_name not in self.lora_config.target_modules: continue base = getattr(self, proj_name) - delta = self._compute_proj_delta(proj_name).to(base.dtype) - base.add_(delta) + delta = self._compute_proj_delta(proj_name) # fp32 + merged = base.to(torch.float32) + delta + base.data.copy_(merged.to(base.dtype)) self.reset_lora_parameters() def forward( 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" From f0649ac6999ad13c0486a9ec2e3c3ca5ccfb16fb Mon Sep 17 00:00:00 2001 From: Ashwinee Panda Date: Thu, 30 Apr 2026 14:49:09 -0700 Subject: [PATCH 16/49] fix: Muon/MoE backward perf regression on Qwen3.5-style MoE * Fix MoE EP backward perf regression for train_router=False Two related changes that recover lost throughput on Qwen3.5-style MoE training when train_router=False (the default for ep_dispatch='alltoall' and the only supported setting for ep_dispatch='deepep'): 1. Triton/Quack EP group GEMM backward (src/xorl/ops/moe/triton.py:127, src/xorl/ops/moe/quack.py:316): skip the extra full down-GEMM that computes grad_expert_scores when expert_scores does not require a gradient. With train_router=False MoEBlock detaches routing_weights upstream so ctx.needs_input_grad[5] is False, making the GEMM purely wasted work in backward. train_router=True still computes grads normally. 2. Make routing replay's record/pop of routing_weights opt-in via a new model arg record_routing_weights (default True for safety). When attention forward is deterministic across checkpoint recompute, the regathered routing_weights match the recorded ones, so the cache is unnecessary and disabling it avoids per-MoE-layer pinned CPU allocations + D2H/H2D copies on every step. The arg is threaded through ModelArguments / ServerArguments -> Trainer / ModelRunner -> build_training_model -> build_foundation_model -> config.record_routing_weights -> MoEBlock.from_config. * Switch Muon ns_algorithm default to gram_newton_schulz + batch standard NS The standard_newton_schulz path was the dominant cause of a ~2x training-step regression on Qwen3.5-35B-A3B vs the prior baseline. With grad shape [E, H, I] (E=local experts) it ran an O(E)-deep Python list comprehension calling the upstream 2D _zeropower_via_newtonschulz once per matrix, emitting ~14k extra kernel launches per optimizer step (40 MoE layers x 3 weight pieces x 8 experts x 5 NS steps x ~3 matmuls each). Two complementary fixes: 1. Default the Muon Newton-Schulz backend to gram_newton_schulz in arguments / server_arguments / optimizer factory / Muon.__init__. gram_newton_schulz already batches across experts via baddbmm and recovered the full ~2x throughput in a 1-pod test (tok/s 12k -> 27k, matching the prior baseline's 27,651 tok/s). 2. Add _batched_zeropower_via_newtonschulz that batches the upstream standard NS recurrence with bmm/baddbmm for users who pin ns_algorithm='standard_newton_schulz'. Bit-exact vs the per-matrix loop on representative shapes; ~2.76x faster on a single CUDA microbench at [8, 2048, 768]. The existing test_muon_standard_newton_schulz_preserves_batched_leading_dims test was monkeypatching the per-matrix _zeropower_via_newtonschulz; updated it to monkeypatch the batched variant and check the [B, H, I] flattening contract end-to-end. (cherry picked from commit 130620902de06418c2d180f635338a1143d0e05a) * ci: re-trigger workflows after title update --------- Co-authored-by: Qingyang Wu --- src/xorl/arguments.py | 18 ++++++-- src/xorl/models/auto.py | 2 + src/xorl/models/layers/moe/moe_block.py | 38 ++++++++-------- src/xorl/ops/moe/quack.py | 5 ++- src/xorl/ops/moe/triton.py | 5 ++- src/xorl/optim/muon.py | 58 ++++++++++++++++++++++--- src/xorl/optim/optimizer.py | 4 +- src/xorl/server/runner/model_runner.py | 1 + src/xorl/server/server_arguments.py | 18 ++++++-- src/xorl/trainers/model_builder.py | 2 + src/xorl/trainers/trainer.py | 1 + tests/optim/test_muon.py | 15 ++++--- 12 files changed, 128 insertions(+), 39 deletions(-) diff --git a/src/xorl/arguments.py b/src/xorl/arguments.py index 35a1f3ee..6a8324ff 100644 --- a/src/xorl/arguments.py +++ b/src/xorl/arguments.py @@ -455,6 +455,16 @@ class ModelArguments: "Disabled by default and must remain False when ep_dispatch='deepep'." }, ) + 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')."}, @@ -582,10 +592,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( diff --git a/src/xorl/models/auto.py b/src/xorl/models/auto.py index 2015aad4..6adef78b 100644 --- a/src/xorl/models/auto.py +++ b/src/xorl/models/auto.py @@ -110,6 +110,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, @@ -151,6 +152,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 diff --git a/src/xorl/models/layers/moe/moe_block.py b/src/xorl/models/layers/moe/moe_block.py index c33acc90..1bd36e0f 100644 --- a/src/xorl/models/layers/moe/moe_block.py +++ b/src/xorl/models/layers/moe/moe_block.py @@ -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) @@ -182,23 +184,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 +295,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/ops/moe/quack.py b/src/xorl/ops/moe/quack.py index 7ed1de10..f0ed194e 100644 --- a/src/xorl/ops/moe/quack.py +++ b/src/xorl/ops/moe/quack.py @@ -311,8 +311,11 @@ def backward(ctx, grad_output): expert_scores = expert_scores.to(gated_output.dtype) # 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: + 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 ) diff --git a/src/xorl/ops/moe/triton.py b/src/xorl/ops/moe/triton.py index de440c8c..e20cadf9 100644 --- a/src/xorl/ops/moe/triton.py +++ b/src/xorl/ops/moe/triton.py @@ -122,8 +122,11 @@ def backward(ctx, grad_output): expert_scores = expert_scores.to(gated_output.dtype) # 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: + 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 diff --git a/src/xorl/optim/muon.py b/src/xorl/optim/muon.py index 95d13761..37691f6b 100644 --- a/src/xorl/optim/muon.py +++ b/src/xorl/optim/muon.py @@ -47,6 +47,55 @@ GROUPED_GRAM_NS_FP32_BYTE_LIMIT = 2 * 1024**3 +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``. @@ -169,7 +218,7 @@ 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, @@ -450,10 +499,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( diff --git a/src/xorl/optim/optimizer.py b/src/xorl/optim/optimizer.py index 7a1aceae..582075b5 100644 --- a/src/xorl/optim/optimizer.py +++ b/src/xorl/optim/optimizer.py @@ -163,7 +163,7 @@ 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"), @@ -354,7 +354,7 @@ def build_optimizer( - signsgd: no optimizer-specific kwargs - muon: {"muon_lr": 0.02, "muon_momentum": 0.95, "muon_nesterov": True, "muon_ns_steps": 5, "muon_adjust_lr_fn": None, - "muon_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, diff --git a/src/xorl/server/runner/model_runner.py b/src/xorl/server/runner/model_runner.py index 7725947c..f894eda5 100644 --- a/src/xorl/server/runner/model_runner.py +++ b/src/xorl/server/runner/model_runner.py @@ -480,6 +480,7 @@ 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), diff --git a/src/xorl/server/server_arguments.py b/src/xorl/server/server_arguments.py index 5ce6118c..aa310722 100644 --- a/src/xorl/server/server_arguments.py +++ b/src/xorl/server/server_arguments.py @@ -91,6 +91,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')."} ) @@ -281,10 +291,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." }, ) @@ -560,6 +571,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, diff --git a/src/xorl/trainers/model_builder.py b/src/xorl/trainers/model_builder.py index 497fad5b..9c9ffd84 100644 --- a/src/xorl/trainers/model_builder.py +++ b/src/xorl/trainers/model_builder.py @@ -106,6 +106,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, @@ -174,6 +175,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, diff --git a/src/xorl/trainers/trainer.py b/src/xorl/trainers/trainer.py index 3fd79905..bc136ded 100644 --- a/src/xorl/trainers/trainer.py +++ b/src/xorl/trainers/trainer.py @@ -402,6 +402,7 @@ 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, diff --git a/tests/optim/test_muon.py b/tests/optim/test_muon.py index 0b3ff118..f3cf1215 100644 --- a/tests/optim/test_muon.py +++ b/tests/optim/test_muon.py @@ -344,13 +344,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 +362,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]], From 462b6840a5f6ea46b5128734c959789b0e7cccc5 Mon Sep 17 00:00:00 2001 From: Ashwinee Panda Date: Thu, 30 Apr 2026 15:12:56 -0700 Subject: [PATCH 17/49] Add Cautious Weight Decay (CWD) * Add Cautious Weight Decay (CWD) Implement Cautious Weight Decay from Chen et al. (arXiv:2510.12402): mask the decoupled weight-decay term by I(u_t * x_t >= 0) so decay only acts on coordinates whose optimizer update aligns with the parameter. The original objective is preserved (no implicit regularizer), and the modification is a one-line change with no extra hyperparameters. Plumbed via a top-level `cautious_weight_decay` flag on TrainArgs and ServerTrainArgs. Supported for anyprecision_adamw, signsgd, muon, and adamw (which auto-routes to AnyPrecisionAdamW with fp32 state since the fused torch.optim.AdamW kernel has no per-coordinate decay hook). SGD explicitly rejects the flag. Mask sign proxy chosen per optimizer: - AdamW family: exp_avg (denominator is positive, so sign matches u_t) - SignSGD: grad - Muon: post-Newton-Schulz update tensor (the actual u_t) Tests cover the mask helper, each optimizer's masked path, the cautious==standard equivalence when all signs align, the build_optimizer routing for adamw/signsgd/anyprecision_adamw/muon/sgd, and the post-NS-update mask for Muon. * Address review: rename Muon test to assert hand-computed reference --- src/xorl/arguments.py | 11 + src/xorl/optim/anyprecision_adamw.py | 25 +- src/xorl/optim/cautious.py | 47 +++ src/xorl/optim/muon.py | 30 +- src/xorl/optim/optimizer.py | 53 +++- src/xorl/optim/signsgd.py | 21 +- src/xorl/server/runner/model_runner.py | 1 + src/xorl/server/server_arguments.py | 10 + src/xorl/trainers/trainer.py | 1 + tests/optim/test_cautious_weight_decay.py | 369 ++++++++++++++++++++++ 10 files changed, 555 insertions(+), 13 deletions(-) create mode 100644 src/xorl/optim/cautious.py create mode 100644 tests/optim/test_cautious_weight_decay.py diff --git a/src/xorl/arguments.py b/src/xorl/arguments.py index 6a8324ff..3522308f 100644 --- a/src/xorl/arguments.py +++ b/src/xorl/arguments.py @@ -545,6 +545,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."}, 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/muon.py b/src/xorl/optim/muon.py index 37691f6b..1ecf4c51 100644 --- a/src/xorl/optim/muon.py +++ b/src/xorl/optim/muon.py @@ -39,6 +39,7 @@ 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 @@ -223,6 +224,7 @@ def __init__( gram_newton_schulz_num_restarts: int = 1, gram_newton_schulz_restart_iterations: Optional[Iterable[int]] = None, 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"}: @@ -270,6 +272,7 @@ def __init__( use_muon=True, adamw_betas=adamw_betas, adamw_eps=adamw_eps, + cautious=cautious, ) Optimizer.__init__(self, params, defaults) @@ -318,6 +321,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) @@ -479,8 +483,15 @@ def _muon_step(self, group: dict) -> None: # 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) @@ -597,6 +608,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: @@ -618,14 +630,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 582075b5..ee0e913a 100644 --- a/src/xorl/optim/optimizer.py +++ b/src/xorl/optim/optimizer.py @@ -114,11 +114,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,16 +171,23 @@ 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 == "muon": adamw_state_dtype = _ANYPRECISION_STATE_DTYPES.get(optimizer_dtype) @@ -174,6 +212,7 @@ 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 @@ -194,6 +233,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. @@ -214,6 +254,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) @@ -333,6 +374,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. @@ -360,6 +402,11 @@ def build_optimizer( "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). """ # 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): @@ -376,6 +423,7 @@ def build_optimizer( no_decay_modules, no_decay_params, optimizer_kwargs=optimizer_kwargs, + cautious_weight_decay=cautious_weight_decay, ) kwargs = optimizer_kwargs or {} @@ -427,6 +475,7 @@ def build_optimizer( fused=fused, optimizer_dtype=optimizer_dtype, optimizer_kwargs=optimizer_kwargs, + cautious_weight_decay=cautious_weight_decay, ) @@ -443,6 +492,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 @@ -526,6 +576,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/server/runner/model_runner.py b/src/xorl/server/runner/model_runner.py index f894eda5..9b4d580f 100644 --- a/src/xorl/server/runner/model_runner.py +++ b/src/xorl/server/runner/model_runner.py @@ -591,6 +591,7 @@ def _initialize_optimizer(self): optimizer_type=optimizer_type, optimizer_dtype=self.train_config.get("optimizer_dtype", "bf16"), optimizer_kwargs=optimizer_kwargs, + cautious_weight_decay=self.train_config.get("cautious_weight_decay", False), ) # Register optimizer pre-hook if available diff --git a/src/xorl/server/server_arguments.py b/src/xorl/server/server_arguments.py index aa310722..9579c1e4 100644 --- a/src/xorl/server/server_arguments.py +++ b/src/xorl/server/server_arguments.py @@ -260,6 +260,15 @@ class ServerArguments: 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={ @@ -610,6 +619,7 @@ def to_config_dict(self) -> Dict[str, Any]: "ce_mode": self.ce_mode, "optimizer": self.optimizer, "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, diff --git a/src/xorl/trainers/trainer.py b/src/xorl/trainers/trainer.py index bc136ded..b20e5917 100644 --- a/src/xorl/trainers/trainer.py +++ b/src/xorl/trainers/trainer.py @@ -611,6 +611,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) 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) From e1e4f905c2150ee6f8781c30d3bd99349cfdaf1e Mon Sep 17 00:00:00 2001 From: Conner Manuel <57027354+connermanuel@users.noreply.github.com> Date: Thu, 30 Apr 2026 15:57:42 -0700 Subject: [PATCH 18/49] fix: resolve dp_sp alias in loss_group/loss_mesh Co-authored-by: Qingyang Wu --- src/xorl/distributed/parallel_state.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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: From 83b37f04426d7f593447a5764a5ec66824437570 Mon Sep 17 00:00:00 2001 From: Qingyang Wu Date: Fri, 1 May 2026 09:29:27 -0700 Subject: [PATCH 19/49] chore(lora): remove unused stacked LoRA helpers, drop lora_utils.py MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * chore(lora): remove unused merge_lora_weights_stacked / unmerge_lora_weights_stacked These two helpers in src/xorl/ops/group_gemm/kernel/lora_utils.py have no in-tree callers and are only re-exported through xorl.lora and xorl.ops.group_gemm.kernel. They date back to the original MoE+LoRA+EP support commit (74be866) and were superseded by MoEExpertsLoRA.merge_weights (the path PR actually patched for fp32-cast-once precision). Removes the function definitions and prunes them from both __init__.py re-export lists. get_lora_delta_weight_stacked is also unused but is left in place for now; flag separately if it should go too. * chore(lora): remove unused stacked LoRA helpers Removes four dead helpers from src/xorl/ops/group_gemm/kernel/lora_utils.py: - merge_lora_weights_stacked - unmerge_lora_weights_stacked - init_lora_weights_stacked - get_lora_delta_weight_stacked None have any in-tree callers (verified by grep across.py /.yaml /.yml / .toml /.json /.md plus check for "import *" forms). They've been dead since the original MoE+LoRA+EP commit (74be866) and were superseded by MoEExpertsLoRA.merge_weights / per-module LoRA paths. The remaining helper, compute_lora_scaling, is kept β€” it is used by src/xorl/models/layers/moe/lora.py and src/xorl/qlora/modules/moe_experts.py. Both __init__.py re-export lists are pruned accordingly. Risk: external consumers (e.g. SGLang LoRA export, internal tooling) that import these symbols from xorl.lora or xorl.ops.group_gemm.kernel will break their import. Worth a sibling-repo grep before merging. * chore(lora): remove unused stacked LoRA helpers, fold compute_lora_scaling into kernel __init__ Removes four dead helpers from src/xorl/ops/group_gemm/kernel/lora_utils.py: - merge_lora_weights_stacked - unmerge_lora_weights_stacked - init_lora_weights_stacked - get_lora_delta_weight_stacked Verified by repo-wide grep across.py /.ipynb /.yaml /.toml /.json / .md /.sh /.txt plus dynamic-ref and `import *` checks: none of the four have any in-tree caller. They've been dead since the original MoE+LoRA+EP commit (74be866); the live MoE-LoRA path is MoEExpertsLoRA.merge_weights, which is the path actually patched. The only remaining helper, compute_lora_scaling, is folded directly into xorl/ops/group_gemm/kernel/__init__.py and lora_utils.py is deleted β€” keeping a one-function file for a 5-line scaling helper isn't worth it. Updates the one direct submodule importer (xorl.qlora.modules.moe_experts) to import from the package instead of the now-deleted submodule. The other caller (xorl.models.layers.moe.lora) already imported from the package and needs no change. Risk: external consumers (SGLang LoRA export path, internal tooling) that import any of the four removed symbols from xorl.lora or xorl.ops.group_gemm.kernel will break their import. Cross-repo grep recommended before merging. --------- Co-authored-by: Ashwinee Panda --- src/xorl/lora/__init__.py | 12 +- src/xorl/ops/group_gemm/kernel/__init__.py | 31 ++-- src/xorl/ops/group_gemm/kernel/lora_utils.py | 148 ------------------- src/xorl/qlora/modules/moe_experts.py | 2 +- 4 files changed, 20 insertions(+), 173 deletions(-) delete mode 100644 src/xorl/ops/group_gemm/kernel/lora_utils.py 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/ops/group_gemm/kernel/__init__.py b/src/xorl/ops/group_gemm/kernel/__init__.py index c4fe4487..e7fbf583 100644 --- a/src/xorl/ops/group_gemm/kernel/__init__.py +++ b/src/xorl/ops/group_gemm/kernel/__init__.py @@ -1,15 +1,8 @@ +import math + # Group GEMM kernels from .group_gemm import group_gemm_same_mn, group_gemm_same_nk -# LoRA utilities -from .lora_utils import ( - compute_lora_scaling, - get_lora_delta_weight_stacked, - init_lora_weights_stacked, - merge_lora_weights_stacked, - unmerge_lora_weights_stacked, -) - # MoE operations from .moe import ( expert_histogram, @@ -21,6 +14,22 @@ from .quack import quack_group_gemm_same_mn, quack_group_gemm_same_nk +def compute_lora_scaling(lora_alpha: int, r: int, use_rslora: bool = False) -> float: + """Compute the LoRA scaling factor. + + Args: + lora_alpha: LoRA alpha parameter. + r: LoRA rank. + use_rslora: Whether to use rank-stabilized LoRA scaling. + + Returns: + Scaling factor. + """ + if use_rslora: + return lora_alpha / math.sqrt(r) + return lora_alpha / r + + __all__ = [ # Group GEMM "group_gemm_same_mn", @@ -34,9 +43,5 @@ "moe_index_compute", "moe_scatter", # LoRA utilities - "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/ops/group_gemm/kernel/lora_utils.py b/src/xorl/ops/group_gemm/kernel/lora_utils.py deleted file mode 100644 index b6c6a354..00000000 --- a/src/xorl/ops/group_gemm/kernel/lora_utils.py +++ /dev/null @@ -1,148 +0,0 @@ -"""LoRA utilities for MoE implementation. - -This module provides utilities for initializing and managing LoRA weights -in the stacked tensor format used by group GEMM kernels. -""" - -import math -from typing import Optional, Tuple - -import torch -import torch.nn as nn - - -def init_lora_weights_stacked( - num_experts: int, - r: int, - in_features: int, - out_features: int, - init_method: str = "kaiming", - dtype: torch.dtype = torch.float32, - device: Optional[torch.device] = None, -) -> Tuple[torch.Tensor, torch.Tensor]: - """Initialize stacked LoRA weights for all experts. - - Creates lora_A and lora_B tensors with appropriate initialization: - - lora_A: Kaiming uniform or Gaussian initialization - - lora_B: Zero initialization (ensures delta_W = 0 at start) - - Args: - num_experts: Number of experts - r: LoRA rank - in_features: Input feature dimension - out_features: Output feature dimension - init_method: Initialization method ("kaiming" or "gaussian") - dtype: Data type for the tensors - device: Device for the tensors - - 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: 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) - - # 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) - - # Initialize lora_A - if init_method == "kaiming": - for i in range(num_experts): - # Initialize each expert's lora_A with kaiming uniform - nn.init.kaiming_uniform_(lora_A[i], a=math.sqrt(5)) - elif init_method == "gaussian": - nn.init.normal_(lora_A, std=1.0 / r) - else: - raise ValueError(f"Unknown init_method: {init_method}") - - # lora_B is already zeros - - return lora_A, lora_B - - -def compute_lora_scaling(lora_alpha: int, r: int, use_rslora: bool = False) -> float: - """Compute the LoRA scaling factor. - - Args: - lora_alpha: LoRA alpha parameter - r: LoRA rank - use_rslora: Whether to use rank-stabilized LoRA scaling - - Returns: - Scaling factor - """ - if use_rslora: - return lora_alpha / math.sqrt(r) - else: - return lora_alpha / r - - -def merge_lora_weights_stacked( - base_weight: torch.Tensor, - lora_A: torch.Tensor, - lora_B: torch.Tensor, - scaling: float, -) -> torch.Tensor: - """Merge LoRA weights into base weights. - - Computes: W' = W + B @ A * 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] - scaling: LoRA scaling factor - - Returns: - Merged weight tensor [num_experts, out_features, in_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 - return base_weight + delta_weight - - -def unmerge_lora_weights_stacked( - merged_weight: torch.Tensor, - lora_A: torch.Tensor, - lora_B: torch.Tensor, - scaling: float, -) -> torch.Tensor: - """Unmerge LoRA weights from merged weights. - - Computes: W = W' - B @ A * 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] - scaling: LoRA scaling factor - - Returns: - Base weight tensor [num_experts, out_features, in_features] - """ - delta_weight = torch.bmm(lora_B, lora_A) * scaling - return merged_weight - delta_weight - - -def get_lora_delta_weight_stacked( - lora_A: torch.Tensor, - lora_B: torch.Tensor, - scaling: float, -) -> torch.Tensor: - """Compute the LoRA weight delta. - - Computes: delta_W = B @ A * scaling - - Args: - lora_A: LoRA A tensor [num_experts, r, in_features] - lora_B: LoRA B tensor [num_experts, out_features, r] - scaling: LoRA scaling factor - - Returns: - Delta weight tensor [num_experts, out_features, in_features] - """ - return torch.bmm(lora_B, lora_A) * scaling diff --git a/src/xorl/qlora/modules/moe_experts.py b/src/xorl/qlora/modules/moe_experts.py index 2edf0e14..6409b5b0 100644 --- a/src/xorl/qlora/modules/moe_experts.py +++ b/src/xorl/qlora/modules/moe_experts.py @@ -32,7 +32,7 @@ from transformers.utils import cached_file from xorl.lora.modules.base import LoraModule -from xorl.ops.group_gemm.kernel.lora_utils import compute_lora_scaling +from xorl.ops.group_gemm.kernel import compute_lora_scaling from xorl.ops.quantize import ( block_fp8_dequantize_gkn, block_fp8_quantize_gkn, From 54a444282223b81abc4fed8feb7fa4735d69ff4a Mon Sep 17 00:00:00 2001 From: Ashwinee Panda Date: Fri, 1 May 2026 15:16:12 -0700 Subject: [PATCH 20/49] Fix nccl_broadcast TCPStore master deadlock Master TCPStore was created with the PyTorch default wait_for_workers=True, which blocks the constructor until world_size-1 workers connect. The design in init_nccl_group is to create a non-blocking listening master first, then fire /init_weights_update_group at the inference endpoints (workers) from a background thread, then complete the NCCL rendezvous in the main thread via _new_process_group_helper. With the default, the master constructor blocks before the background thread can even start, so workers never receive the HTTP request and sync_inference_weights deadlocks until the request times out. Pass wait_for_workers=False so the master returns as soon as it is listening; the actual rendezvous still synchronizes via _init_training_process_group. Update the unit-test fake TCPStore to accept extra kwargs. --- src/xorl/server/weight_sync/backends/nccl_broadcast.py | 8 ++++++++ tests/server/weight_sync/test_nccl_broadcast_lifecycle.py | 2 +- 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/src/xorl/server/weight_sync/backends/nccl_broadcast.py b/src/xorl/server/weight_sync/backends/nccl_broadcast.py index ec4f3f65..ebd2962f 100644 --- a/src/xorl/server/weight_sync/backends/nccl_broadcast.py +++ b/src/xorl/server/weight_sync/backends/nccl_broadcast.py @@ -155,12 +155,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 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 From b5d0248fd366206274570a5453fc95283e1ed333 Mon Sep 17 00:00:00 2001 From: Ashwinee Panda Date: Mon, 4 May 2026 17:05:16 -0700 Subject: [PATCH 21/49] feat: add OLMo-2 support MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Add OLMo-2 model support Adds Olmo2ForCausalLM under xorl.models.transformers.olmo2, mirroring the Llama 3 / Qwen 3 modules. OLMo-2 differs from Llama in two ways: * Post-norm: layer norms are applied after attention and MLP (post_attention_layernorm, post_feedforward_layernorm); there is no input_layernorm. * Full-axis QK norm: q_norm and k_norm normalize across the entire (num_heads * head_dim) axis prior to reshape, not per-head. Wires the model_type "olmo2" through auto.py so HF configs are loaded into our local Olmo2Config (which carries the TP/PP plans), and adds a checkpoint handler that fuses gate_proj/up_proj into gate_up_proj and q/k/v_proj into qkv_proj on load (and the inverse on save). Includes parity tests against transformers.models.olmo2. * Add tensor-parallel plan for OLMo-2 Wires up TP for OLMo-2 via the standard colwise/rowwise pair, with ``LocalAxisRMSNormShard`` handling the model's full-axis q_norm/k_norm under colwise q/k_proj. Keeps the residual stream Replicate across decoder layers so xorl's loss path (vocab_parallel_cross_entropy matmuls hidden_states @ lm_head.weight.t() directly, bypassing lm_head) sees a full [B, S, H] per rank. * embed_tokens: "embedding" β€” Replicate output (default). * q/k/v_proj, gate/up_proj: ColwiseParallel() β€” Replicate input, Shard(-1) output, use_local_output=True. Block internals run on plain local tensors with hidden/tp per rank. * o_proj/down_proj: RowwiseParallel() β€” Shard(-1) input, Replicate output (all-reduce). Post-norms and residuals see Replicate without any extra plumbing. * q_norm/k_norm: LocalAxisRMSNormShard β€” weight sharded on dim 0 so each rank's slice matches its colwise q/k_proj output. Computes a local-axis RMS, matching HuggingFace's OLMo-2 reference. * lm_head: ColwiseParallel() β€” Replicate input, vocab-parallel output. Add tests/distributed/test_olmo2_tp_e2e.py: 2-rank gloo+CPU end-to-end forward + backward through every TP boundary, including lm_head and gradient-flow assertion. Add tests/models/test_olmo2_support.py::test_olmo2_tp_plan_uses_local_axis_qk_norm: structural assertions on the plan (LocalAxisRMSNormShard for q_norm/k_norm, post-norms + final norm not in the plan, vanilla colwise/rowwise elsewhere). Depends on the generic TP infrastructure in (LocalAxisRMSNormShard, DTensor-aware RMSNorm, ParallelStyle resolver passthrough). H100 1-node smoke (8Γ—H100, v_tp2_dp4.yaml): - 15/15 steps complete - ~72k tok/s steady (peaks 73k) - 40.97 GB peak VRAM - loss converging, gradients healthy Closes * Add tensor-parallel plan for OLMo-2 OLMo-2 declares full-axis q_norm/k_norm (over num_heads * head_dim, not per-head). Under colwise q_proj/k_proj, the q/k tensors arrive hidden-sharded, and a full-hidden RMSNorm weight can't be applied directly β€” so a custom ParallelStyle is needed. Every other model in the repo either has no QK norm or uses per-head + reshape-first, both of which compose with stock ColwiseParallel; the OLMo-2 quirk is scoped to the olmo2/ folder rather than presented as generic TP infra. Pieces: * src/xorl/models/transformers/olmo2/tp_styles.py β€” new LocalAxisRMSNormShard ParallelStyle: shards a 1-D RMSNorm weight along dim 0 with no input/output redistribute. Each rank's weight slice matches its local q/k slice from colwise q/k_proj. * src/xorl/models/transformers/olmo2/modeling_olmo2.py β€” new Olmo2QKRMSNorm subclass of RMSNorm. Without TP the parent forward runs unchanged. With TP the forward detects the Shard(0) DTensor weight and runs the fused op directly on local tensors, computing a local-axis RMS that matches HuggingFace's Olmo2RMSNorm reference under TP. Used only for q_norm/k_norm; post-norms and the final norm stay on vanilla RMSNorm. * src/xorl/models/transformers/olmo2/parallelize.py β€” Replicate- throughout TP plan (default colwise/rowwise everywhere, post-norms and final norm not in the plan). lm_head: ColwiseParallel for use with vocab_parallel_cross_entropy. * src/xorl/distributed/torch_parallelize.py β€” minimal 3-line ParallelStyle passthrough in _resolve_tp_style so plans can declare style instances directly. Backwards-compatible with string-based plans. Tests (CPU + 2-rank gloo, no GPU required): * tests/distributed/test_olmo2_qk_rms_norm.py β€” Olmo2QKRMSNorm forward with and without TP, validated against a per-rank local-slice RMSNorm reference under LocalAxisRMSNormShard. * tests/distributed/test_olmo2_tp_e2e.py β€” full OLMo-2 fwd+bwd through every TP boundary including lm_head and gradient flow. * tests/models/test_olmo2_support.py β€” TP plan structural assertions (LocalAxisRMSNormShard for q_norm/k_norm, post-norms not in plan, stock colwise/rowwise elsewhere). H100 1-node smoke (8Γ—H100, v_tp2_dp4.yaml β€” TP=2, FSDP=4, mb=2): - 15/15 steps complete - ~72k tok/s steady, 40.97 GB peak VRAM - loss converging, gradients healthy Closes, --------- Co-authored-by: Qingyang Wu --- src/xorl/distributed/torch_parallelize.py | 23 +- src/xorl/models/auto.py | 5 + .../models/transformers/olmo2/__init__.py | 0 .../transformers/olmo2/checkpoint_handler.py | 170 +++++++ .../transformers/olmo2/configuration_olmo2.py | 84 ++++ .../transformers/olmo2/modeling_olmo2.py | 424 ++++++++++++++++++ .../models/transformers/olmo2/parallelize.py | 80 ++++ .../models/transformers/olmo2/tp_styles.py | 37 ++ tests/distributed/test_olmo2_qk_rms_norm.py | 112 +++++ tests/distributed/test_olmo2_tp_e2e.py | 137 ++++++ tests/models/test_olmo2_support.py | 180 ++++++++ 11 files changed, 1245 insertions(+), 7 deletions(-) create mode 100644 src/xorl/models/transformers/olmo2/__init__.py create mode 100644 src/xorl/models/transformers/olmo2/checkpoint_handler.py create mode 100644 src/xorl/models/transformers/olmo2/configuration_olmo2.py create mode 100644 src/xorl/models/transformers/olmo2/modeling_olmo2.py create mode 100644 src/xorl/models/transformers/olmo2/parallelize.py create mode 100644 src/xorl/models/transformers/olmo2/tp_styles.py create mode 100644 tests/distributed/test_olmo2_qk_rms_norm.py create mode 100644 tests/distributed/test_olmo2_tp_e2e.py create mode 100644 tests/models/test_olmo2_support.py diff --git a/src/xorl/distributed/torch_parallelize.py b/src/xorl/distributed/torch_parallelize.py index 2c63acbd..70ed5e61 100644 --- a/src/xorl/distributed/torch_parallelize.py +++ b/src/xorl/distributed/torch_parallelize.py @@ -29,6 +29,7 @@ from torch.distributed._composable.fsdp import MixedPrecisionPolicy, fully_shard from torch.distributed.tensor.parallel import ( ColwiseParallel, + ParallelStyle, RowwiseParallel, parallelize_module, ) @@ -71,19 +72,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( diff --git a/src/xorl/models/auto.py b/src/xorl/models/auto.py index 6adef78b..7c8f3207 100644 --- a/src/xorl/models/auto.py +++ b/src/xorl/models/auto.py @@ -56,6 +56,11 @@ def _load_local_xorl_config( return Qwen2Config(**{k: v for k, v in config_dict.items() if not k.startswith("_")}) + if model_type == "olmo2": + from .transformers.olmo2.configuration_olmo2 import Olmo2Config + + return Olmo2Config(**{k: v for k, v in config_dict.items() if not k.startswith("_")}) + return None 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/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/models/test_olmo2_support.py b/tests/models/test_olmo2_support.py new file mode 100644 index 00000000..5ea04789 --- /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=5e-5, rtol=5e-5) + torch.testing.assert_close(xorl_logits, hf_logits, atol=5e-5, rtol=5e-5) From 5a10044d7293c56dc98b89c80a5b09f37e5139a5 Mon Sep 17 00:00:00 2001 From: Qingyang Wu Date: Tue, 5 May 2026 13:52:59 -0700 Subject: [PATCH 22/49] Move vendored xorl.ops.quack ops to xorl_quack:: torch.library namespace MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit flash_attn-cute pulls in upstream ``quack-kernels`` via ``flash_attn.cute.utils`` during normal training-stack imports. Once ``quack-kernels >= 0.3.11`` is installed, both upstream ``quack`` and the vendored ``xorl.ops.quack`` register ``@torch.library.custom_op`` entries in the **same** ``quack::`` namespace β€” e.g. both define ``quack::gemm_out`` with mismatched schemas. Whichever module gets imported second silently wins the registration, which produces ``ValueError: vector::reserve`` deep inside the dispatcher when xorl's quack MoE backend invokes the Python wrapper for the schema it expected. This rename moves all vendored op names to ``xorl_quack::`` (gemm_out, gemm_gated_out, gemm_dgated_out, gemm_act_out, gemm_dact_out, gemm_add_out, gemm_add_inplace, gemm_symmetric_out, _softmax_fwd, _softmax_backward, _topk_fwd, _topk_bwd, _rmsnorm_fwd, _rmsnorm_bwd, cross_entropy_fwd_out, cross_entropy_bwd_out). Vendored and upstream now coexist cleanly: ``torch.ops.xorl_quack.*`` for the local copy, ``torch.ops.quack.*`` for the upstream. No public Python API changes β€” all callers go through ``xorl.ops.quack.gemm_interface.gemm`` etc. Co-authored-by: Ashwinee Panda --- src/xorl/ops/quack/cross_entropy.py | 4 ++-- src/xorl/ops/quack/gemm_interface.py | 22 +++++++++++----------- src/xorl/ops/quack/rmsnorm.py | 4 ++-- src/xorl/ops/quack/softmax.py | 4 ++-- src/xorl/ops/quack/topk.py | 4 ++-- 5 files changed, 19 insertions(+), 19 deletions(-) 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], From 3b644a5de3d1db07c9d9164aeb56248ba9871302 Mon Sep 17 00:00:00 2001 From: Ashwinee Panda Date: Tue, 5 May 2026 15:59:19 -0700 Subject: [PATCH 23/49] Add DistSignSGD optimizer support MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Add DistSignSGD optimizer DistSignSGD signs gradients before FSDP2 reduction, turning the reduce-scatter into a distributed majority-vote update. Complements the existing per-rank SignSGD by aggregating sign votes across DP workers. Also wires in two small training-loop improvements: - build_wandb_init_settings() reads WANDB_INIT_TIMEOUT env var to extend wandb.init timeout on slow-network runs. - _log_memory_snapshot no longer issues an NCCL all_reduce during startup; the snapshot is rank-0 observability only, and the extra collective could wedge setup before step 1 when ranks progressed unevenly. * Address PR review feedback - count_active_microbatches: build per-mb flags into a single tensor and issue one batched all-reduce (op=SUM) instead of an all-reduce +.item() per micro-batch. - get_distsign_grad_scale_factor: divide by the actual number of (mb, rank) pairs that cast a sign vote rather than active_microbatches * dp_size, so ranks with zero valid tokens (sign(0) = 0 abstains) don't bias the per-step update toward zero. - distsignsgd: harden the FSDP2 internals walk and assert no FSDP-managed parameter ends up with the local sign hook installed, preventing silent double-signing if torch FSDP2 internals shift. - Revert unrelated changes: _log_memory_snapshot rank-0/no-collective rewrite and WANDB_INIT_TIMEOUT helper. These are not part of DistSignSGD; ship separately if needed. * Address second-round review feedback on PR - DistSignReduceScatter: force ReduceOp.SUM regardless of the op FSDP passes through. With reduce_dtype=fp32, FSDP2 may pass AVG (divide by N), and the trainer's voter-total scale factor already divides by active_voter_total β€” inheriting AVG would double-divide and silently shrink updates. - sync_sp_gradients: drop the now-unnecessary.to_local() unwrap on DTensor grads. The pre-PR path called dist.all_reduce on p.grad directly; the only behavior we need to add for DistSignSGD is the skip_dtensor_grads early-return. - _register_local_sign_hooks: raise NotImplementedError when a non-FSDP DTensor parameter (e.g. a pure TP/PP shard) shows up. Silently skipping it would let its grads bypass the sign step and pollute the voter total. - configure_distsignsgd: switch to direct ps.{dp_replicate_enabled, ep_enabled, cp_enabled, cp_fsdp_mode, sp_grad_sync_group} access so an attribute rename surfaces as AttributeError instead of silently defaulting to False/None and skipping the guard. Also reject ps.ep_enabled β€” EP-managed grads use a different fsdp group, so a single active_voter_total cannot normalize them correctly. - get_effective_grad_clip_value: clarify in the docstring that under DistSignSGD the value reported as "grad_norm" is really the L2 norm of accumulated sign votes (vote_l2_norm), not a true gradient magnitude. - Tests: cover the AVG β†’ SUM override, EP rejection, and non-FSDP DTensor rejection paths; update sync_sp_gradients tests to match the simplified DTensor handling. --------- Co-authored-by: Qingyang Wu --- src/xorl/arguments.py | 5 +- src/xorl/optim/__init__.py | 2 + src/xorl/optim/distsignsgd.py | 270 ++++++++++++++++++ src/xorl/optim/optimizer.py | 24 +- src/xorl/server/runner/model_runner.py | 65 ++++- src/xorl/server/server_arguments.py | 5 +- src/xorl/trainers/trainer.py | 42 ++- src/xorl/trainers/training_utils.py | 109 ++++++- tests/e2e/qwen3_8b/test_distsignsgd.py | 71 +++++ tests/optim/test_distsignsgd.py | 378 +++++++++++++++++++++++++ tests/server/test_server_arguments.py | 23 ++ tests/test_arguments.py | 34 +++ tests/trainers/test_training_utils.py | 153 ++++++++++ 13 files changed, 1153 insertions(+), 28 deletions(-) create mode 100644 src/xorl/optim/distsignsgd.py create mode 100644 tests/e2e/qwen3_8b/test_distsignsgd.py create mode 100644 tests/optim/test_distsignsgd.py create mode 100644 tests/trainers/test_training_utils.py diff --git a/src/xorl/arguments.py b/src/xorl/arguments.py index 3522308f..e0b96e58 100644 --- a/src/xorl/arguments.py +++ b/src/xorl/arguments.py @@ -565,10 +565,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." }, ) 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/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/optimizer.py b/src/xorl/optim/optimizer.py index ee0e913a..5942f327 100644 --- a/src/xorl/optim/optimizer.py +++ b/src/xorl/optim/optimizer.py @@ -12,6 +12,7 @@ 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 .signsgd import SignSGD @@ -189,6 +190,9 @@ def _get_optimizer_cls_and_kwargs( elif optimizer_type == "signsgd": 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") @@ -218,7 +222,7 @@ def _get_optimizer_cls_and_kwargs( return Muon, ctor_kwargs else: raise ValueError( - f"Unsupported optimizer type: '{optimizer_type}'. Supported: adamw, anyprecision_adamw, sgd, signsgd, muon." + f"Unsupported optimizer type: '{optimizer_type}'. Supported: adamw, anyprecision_adamw, sgd, signsgd, distsignsgd, muon." ) @@ -242,6 +246,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 """ @@ -386,7 +391,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. @@ -394,6 +399,7 @@ 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": "gram_newton_schulz", @@ -408,9 +414,13 @@ def build_optimizer( 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, @@ -425,6 +435,9 @@ def build_optimizer( optimizer_kwargs=optimizer_kwargs, cautious_weight_decay=cautious_weight_decay, ) + if optimizer_type == "distsignsgd": + configure_distsignsgd(model) + return optimizer kwargs = optimizer_kwargs or {} @@ -465,7 +478,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, @@ -477,6 +490,9 @@ def build_optimizer( optimizer_kwargs=optimizer_kwargs, cautious_weight_decay=cautious_weight_decay, ) + if optimizer_type == "distsignsgd": + configure_distsignsgd(model) + return optimizer def build_ep_fsdp2_optimizer( diff --git a/src/xorl/server/runner/model_runner.py b/src/xorl/server/runner/model_runner.py index 9b4d580f..360c6204 100644 --- a/src/xorl/server/runner/model_runner.py +++ b/src/xorl/server/runner/model_runner.py @@ -47,11 +47,15 @@ ) 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, negotiate_pp_seq_len, pad_micro_batches_for_pp, pp_loss_fn, + scale_model_gradients, sync_sp_gradients, ) from xorl.trainers.training_utils import ( @@ -217,6 +221,8 @@ 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] = {} @@ -562,6 +568,9 @@ def _resolve_lora_target_modules(self) -> List[str]: def _initialize_optimizer(self): """Initialize the optimizer.""" optimizer_type = self.train_config.get("optimizer", "adamw") + 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 = None if optimizer_type == "muon": optimizer_kwargs = { @@ -770,6 +779,11 @@ 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) + # ========================================================================= # Loss computation dispatch # ========================================================================= @@ -1000,6 +1014,10 @@ def _forward_loop( # Count valid tokens globally global_valid_tokens = self._count_global_valid_tokens(micro_batches) + if compute_backward and self._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") @@ -1137,11 +1155,22 @@ 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=self._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 self._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 + ) # Build result result = { @@ -1387,6 +1416,10 @@ def forward_backward( # PP path if self.pp_enabled: global_valid_tokens = self._count_global_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 # 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): @@ -1407,6 +1440,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 self._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() @@ -1564,6 +1604,8 @@ 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 = self._accumulated_active_microbatches.pop(model_id, 0) + accumulated_active_voter_total = self._accumulated_active_voter_total.pop(model_id, 0) # Multi-adapter path: use adapter's own optimizer on adapter's own parameters if self._adapter_manager is not None: @@ -1587,13 +1629,14 @@ def optim_step( # 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 self._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: @@ -1602,6 +1645,10 @@ def optim_step( param_group["lr"] = effective_lr ps = get_parallel_state() + clip_value = get_effective_grad_clip_value( + clip_value, + use_distsignsgd=self._use_distsignsgd, + ) grad_norm = clip_gradients( self.model, @@ -1646,6 +1693,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" ) diff --git a/src/xorl/server/server_arguments.py b/src/xorl/server/server_arguments.py index 9579c1e4..a3df5ab8 100644 --- a/src/xorl/server/server_arguments.py +++ b/src/xorl/server/server_arguments.py @@ -247,10 +247,11 @@ 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." }, ) diff --git a/src/xorl/trainers/trainer.py b/src/xorl/trainers/trainer.py index b20e5917..bd4b5dcf 100644 --- a/src/xorl/trainers/trainer.py +++ b/src/xorl/trainers/trainer.py @@ -55,12 +55,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, 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 @@ -603,6 +607,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, @@ -842,6 +847,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: @@ -861,7 +870,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) @@ -879,6 +892,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, @@ -889,7 +909,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.""" @@ -997,12 +1021,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: @@ -1010,9 +1030,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, ) diff --git a/src/xorl/trainers/training_utils.py b/src/xorl/trainers/training_utils.py index 2a85cc50..50880a13 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, 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/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/server/test_server_arguments.py b/tests/server/test_server_arguments.py index 1f63a119..c09914a0 100644 --- a/tests/server/test_server_arguments.py +++ b/tests/server/test_server_arguments.py @@ -30,6 +30,29 @@ def test_load_server_arguments_threads_signsgd_through_nested_config(tmp_path): assert args.to_config_dict()["train"]["optimizer"] == "signsgd" +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_muon_gram_newton_schulz_through_nested_config(tmp_path): config_path = tmp_path / "server_config.yaml" config_path.write_text( diff --git a/tests/test_arguments.py b/tests/test_arguments.py index 07d32bc6..157fef54 100644 --- a/tests/test_arguments.py +++ b/tests/test_arguments.py @@ -44,6 +44,40 @@ def test_parse_args_accepts_signsgd_from_yaml(tmp_path, monkeypatch): assert args.train.optimizer_kwargs == {} +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): config_path = tmp_path / "config.yaml" config_path.write_text( 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" From b5722dae2fefdc076d4ab1811215740a2f2496ae Mon Sep 17 00:00:00 2001 From: Ashwinee Panda Date: Wed, 6 May 2026 13:37:05 -0700 Subject: [PATCH 24/49] fix: preserve FA3 varlen metadata after sync padding * Fix FA3 varlen metadata after sync padding * test: cover stale sync-padded cu seqlens --- src/xorl/distributed/sync_padding.py | 21 ++++++++ tests/distributed/test_sync_padding.py | 71 ++++++++++++++++++++++++++ 2 files changed, 92 insertions(+) create mode 100644 tests/distributed/test_sync_padding.py 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/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) From 04175e2bebb4c69edd2c450a3d0bd2c96390d980 Mon Sep 17 00:00:00 2001 From: Zhongzhu Zhou Date: Fri, 8 May 2026 08:02:26 +1000 Subject: [PATCH 25/49] feat(models): add GLM-4 MoE (GLM-4.5/4.6/4.7) training support MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Add GLM-4 MoE (GLM-4.5/4.6/4.7) training support Port GLM-4 MoE architecture from tomni to xorl, supporting GLM-4.5-Air, GLM-4.5, GLM-4.6, and GLM-4.7 model variants (model_type: glm4_moe). Key architecture features: - Sigmoid-based grouped top-k routing with e_score_correction_bias - Shared experts alongside routed experts - Partial rotary embeddings (partial_rotary_factor) - Optional QK normalization - Dense layers for first_k_dense_replace layers Includes EP/TP parallelization, checkpoint handler for HF weight conversion, and dummy-data benchmark configs (EP=8, EP+Ulysses). Fix _naive_apply_rotary_pos_emb to handle partial rotary automatically when cos/sin dim < head_dim, which was needed for UlyssesAsyncStrategy compatibility with partial-rotary models. * Fix EP plan to use fused gate_up_proj for GLM-4 MoE experts * Fix ruff lint: import sorting and formatting in GLM-4 MoE module Made-with: Cursor * Fix GLM-4.7 MoE routing bug and add model support - Fix e_score_correction_bias routing bug: bias was incorrectly included in routing weights (expert contribution), not just top-k selection. This caused K3 divergence ~10^28 vs SGLang (now K3 β‰ˆ 0). - Add MTP layer weight remapping in checkpoint handler (embed_tokens, norm, lm_head stored under model.layers.92 in HF checkpoint). - Register glm4_moe model type in auto config loader. - Add GLM-4 MoE FLOPs counter with shared expert and dense layer support. - Fix token shifting in non-packing path (packing.py): HF-format labels were not shifted when enable_packing=False. - Minor K3 test script updates for GLM-4.7 compatibility. * Lint cleanup: hoist imports, fix formatting for GLM-4 MoE PR - Hoist `import warnings` to top-level in checkpoint_handler.py - Hoist `get_replay_stage` import to top-level in modeling_glm4_moe.py - Fix ruff format (trailing whitespace alignment) in count_flops.py - Add experiments/ to.gitignore * Fix MoE test failures: propagate moe_act, fused gate_up_proj, skip server tests - base.py: propagate moe_checkpoint_method='moe_act' to MoEExperts modules by setting _moe_act=True during gradient_checkpointing_enable - test_moe_gkn_format.py: pass fused gate_up_proj to TritonMoeExpertsFunction which now requires it for the group_gemm path - conftest.py: auto-skip server-marked e2e tests unless XORL_RUN_SERVER_TESTS=1 to prevent failures in standard pytest runs * Revert packing.py changes from PR The HF-format token shifting fix for SequentialPacker is generic data-pipeline correctness, not GLM-4-specific, and a more complete version of the fix (which also shifts position_ids and handles weights/advantages) is being landed independently on split/trainer-non-packed-shift. * Drop incidental edits from PR:.gitignore, e2e conftest, moe_gkn test, base.py * experiments(k3_tests): raise on out-of-vocab token IDs instead of skipping Silently filtering OOV token IDs masked a real bug β€” usually SGLang TP padding leaking into the output. Raise ValueError with the same diagnostic message so the failing prompt is the failing prompt rather than a silently-dropped sample. * fix(glm4_moe): use local Glm4MoeRMSNorm with hardcoded fp32 upcast The shared RMSNorm class respects the global rmsnorm_mode (default "native" β†’ F.rms_norm in bf16). On GLM-4 MoE that path interacts pathologically with the triton/quack MoE kernels: K3 KL divergence vs SGLang/HF on GLM-4.5-Air showed a single-token logprob ratio of ~23 (mean K3 0.011, max 0.185) on borderline routing decisions, while every other rmsnorm Γ— moe combination held mean K3 < 0.002. Match HF Glm4MoeRMSNorm: hardcoded fp32 variance + bf16 weight multiply. No global-mode dependency. Other models keep using the shared RMSNorm. * ci: trigger checks * perf(glm4_moe): compile the fp32-upcast RMSNorm forward Glm4MoeRMSNorm intentionally bypasses F.rms_norm to keep an fp32 upcast (the bf16 native path compounds error across GLM-4 MoE's 92+ norm layers). The previous implementation paid the per-call cost of ~5 separate kernel launches Γ— 92 layers β‰ˆ 460 launches per forward pass. Adds a compiled_eager_rms_norm helper to xorl.models.layers. normalization (parallel to compiled_rms_norm and compiled_zero_centered_rms_norm) that wraps the existing eager_rms_norm in torch.compile, fusing the cast/pow/mean/rsqrt/ mul chain into a single Inductor kernel. Glm4MoeRMSNorm.forward now delegates to it. Behaviour unchanged; same fp32 numerics. * style: collapse compiled_eager_rms_norm signature for ruff-format --------- Co-authored-by: Qingyang Wu --- .../dummy/configs/full/glm4_moe_cp1_sp8.yaml | 56 ++ .../dummy/configs/full/glm4_moe_ep8.yaml | 54 ++ src/xorl/models/auto.py | 4 + src/xorl/models/layers/normalization.py | 10 + src/xorl/models/layers/rope.py | 12 +- src/xorl/models/transformers/__init__.py | 2 + .../models/transformers/glm4_moe/__init__.py | 38 + .../glm4_moe/checkpoint_handler.py | 345 +++++++++ .../glm4_moe/configuration_glm4_moe.py | 174 +++++ .../glm4_moe/modeling_glm4_moe.py | 721 ++++++++++++++++++ .../transformers/glm4_moe/parallelize.py | 63 ++ src/xorl/utils/count_flops.py | 63 ++ 12 files changed, 1540 insertions(+), 2 deletions(-) create mode 100644 examples/local/dummy/configs/full/glm4_moe_cp1_sp8.yaml create mode 100644 examples/local/dummy/configs/full/glm4_moe_ep8.yaml create mode 100644 src/xorl/models/transformers/glm4_moe/__init__.py create mode 100644 src/xorl/models/transformers/glm4_moe/checkpoint_handler.py create mode 100644 src/xorl/models/transformers/glm4_moe/configuration_glm4_moe.py create mode 100644 src/xorl/models/transformers/glm4_moe/modeling_glm4_moe.py create mode 100644 src/xorl/models/transformers/glm4_moe/parallelize.py 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/src/xorl/models/auto.py b/src/xorl/models/auto.py index 7c8f3207..bee168cc 100644 --- a/src/xorl/models/auto.py +++ b/src/xorl/models/auto.py @@ -16,6 +16,7 @@ from .layers.attention import ATTENTION_FUNCTIONS from .layers.normalization import set_rmsnorm_mode from .loader import ModelLoader, get_loader +from .transformers.glm4_moe.configuration_glm4_moe import Glm4MoeConfig 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 ( @@ -45,6 +46,9 @@ 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)) 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/transformers/__init__.py b/src/xorl/models/transformers/__init__.py index a4bec082..5af9d2d1 100644 --- a/src/xorl/models/transformers/__init__.py +++ b/src/xorl/models/transformers/__init__.py @@ -1,4 +1,5 @@ from . import ( + glm4_moe, llama3, qwen2, qwen3, @@ -9,6 +10,7 @@ __all__ = [ + "glm4_moe", "llama3", "qwen2", "qwen3", 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/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 From 3515add447c9c1cf459cd043bbb01b7198ec1a1d Mon Sep 17 00:00:00 2001 From: Conner Manuel <57027354+connermanuel@users.noreply.github.com> Date: Thu, 7 May 2026 16:39:46 -0700 Subject: [PATCH 26/49] feat: Allow loss funcs to accept reducers * Update loss aggregation to normalize across minibatches and ranks * feat: Add reducers support for PPO and IS * refactor: use reducers in model_runner * add SequenceSum * upd surface of losses to return metric tensors instead of item * cleanup * add durations to pytest * metric reduction tests * mean->partial, remove sum * use generalized cu_seq_lens form for seqpartial * remove closure, rename loss -> local_loss_sum --- pyproject.toml | 4 + src/xorl/ops/loss/__init__.py | 4 + src/xorl/ops/loss/causallm_loss.py | 24 +- src/xorl/ops/loss/grpo_loss.py | 90 +++---- src/xorl/ops/loss/importance_sampling_loss.py | 67 ++--- src/xorl/ops/loss/loss_output.py | 9 +- src/xorl/ops/loss/policy_loss.py | 110 +++++--- src/xorl/ops/loss/reducers.py | 71 +++++ src/xorl/server/runner/model_runner.py | 245 ++++++++++-------- .../test_loss_metric_reductions.py | 225 ++++++++++++++++ tests/ops/loss/test_drgrpo_loss.py | 59 +++-- .../ops/loss/test_importance_sampling_loss.py | 109 ++++++++ tests/ops/loss/test_policy_loss.py | 129 +++++++++ tests/ops/loss/test_reducers.py | 197 ++++++++++++++ 14 files changed, 1071 insertions(+), 272 deletions(-) create mode 100644 src/xorl/ops/loss/reducers.py create mode 100644 tests/distributed/test_loss_metric_reductions.py create mode 100644 tests/ops/loss/test_importance_sampling_loss.py create mode 100644 tests/ops/loss/test_policy_loss.py create mode 100644 tests/ops/loss/test_reducers.py diff --git a/pyproject.toml b/pyproject.toml index 4cb41d74..0abc57a3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -114,6 +114,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/ops/loss/__init__.py b/src/xorl/ops/loss/__init__.py index ec7eed53..501064b8 100644 --- a/src/xorl/ops/loss/__init__.py +++ b/src/xorl/ops/loss/__init__.py @@ -15,6 +15,7 @@ from xorl.ops.loss.importance_sampling_loss import importance_sampling_loss_function from xorl.ops.loss.loss_output import LossOutput from xorl.ops.loss.policy_loss import policy_loss_function +from xorl.ops.loss.reducers import Reducer, SequencePartial, TokenPartial from xorl.ops.loss.vocab_parallel_cross_entropy import vocab_parallel_cross_entropy @@ -51,6 +52,9 @@ def register_loss_function(name: str, fn: Callable) -> None: "CrossEntropyMode", "LossOutput", "LOSS_REGISTRY", + "Reducer", + "SequencePartial", + "TokenPartial", "get_loss_function", "register_loss_function", "causallm_loss_function", diff --git a/src/xorl/ops/loss/causallm_loss.py b/src/xorl/ops/loss/causallm_loss.py index cd9d643f..75c212ff 100644 --- a/src/xorl/ops/loss/causallm_loss.py +++ b/src/xorl/ops/loss/causallm_loss.py @@ -8,6 +8,7 @@ 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 @@ -22,6 +23,7 @@ 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": """ @@ -42,6 +44,12 @@ 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 @@ -64,6 +72,11 @@ 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: @@ -83,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, @@ -93,7 +106,6 @@ def causallm_loss_function( return LossOutput(loss=loss) z_loss_enabled = z_loss_coef > 0.0 - valid_count = valid_mask.sum().clamp(min=1) if return_per_token: # Compute cross-entropy based on mode (and Z-loss when enabled). @@ -117,9 +129,9 @@ def causallm_loss_function( lse = torch.logsumexp(logits_flat, dim=-1) per_token_lse_sq = (lse * lse) * valid_mask.to(lse.dtype) - ce_loss = per_token_ce.sum() / valid_count + ce_loss = loss_reducer(per_token_ce, mask_flat) if z_loss_enabled: - z_loss = per_token_lse_sq.sum() / valid_count + 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: @@ -156,9 +168,9 @@ def causallm_loss_function( lse = torch.logsumexp(logits_flat, dim=-1) per_token_lse_sq = (lse * lse) * valid_mask.to(lse.dtype) - ce_loss = per_token_ce.sum() / valid_count + ce_loss = loss_reducer(per_token_ce, mask_flat) if z_loss_enabled: - z_loss = per_token_lse_sq.sum() / valid_count + 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/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/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/server/runner/model_runner.py b/src/xorl/server/runner/model_runner.py index 360c6204..6c3747c7 100644 --- a/src/xorl/server/runner/model_runner.py +++ b/src/xorl/server/runner/model_runner.py @@ -15,6 +15,7 @@ import gc import logging +import math import os import time from typing import Any, Dict, List, Optional @@ -33,6 +34,7 @@ from xorl.lora import LoraLinear from xorl.models.layers.moe.routing_replay import set_replay_stage from xorl.ops.loss import ( + TokenPartial, causallm_loss_function, importance_sampling_loss_function, policy_loss_function, @@ -90,62 +92,45 @@ 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 _sp_allreduce_kl_metrics( + metrics: Dict[str, Any], + metric_ops: Optional[Dict[str, str]], + sp_group, +) -> Dict[str, Any]: """ All-reduce KL/ratio metrics across the sequence-parallel (Ulysses) group. - With Ulysses SP, each rank only sees a shard of the sequence. Rank 0 often - has only prompt tokens (all target_tokens=-100), so its KL stats are zeros. - 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. + With Ulysses SP, each rank sees only a shard of the sequence. Rank 0 often + has only prompt tokens (target_tokens=-100), so its KL stats are zeros. + Reducing here makes per-mb metrics CP-global before downstream cross-mb / + cross-DP accumulation. Means are raw partial sums; min/max use Β±inf + identity. Nothing is finalized β€” downstream sum-then-divide still applies. """ + metric_ops = metric_ops or {} + if not metrics: + return metrics + 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) - - # --- Update metrics with properly reduced values --- - for key in mean_keys: - if key in sum_tensors: - metrics[key] = sum_tensors[key].item() / total_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 - - # Clean up internal key - metrics.pop("_n_valid_kl", None) + by_op: Dict[str, list[str]] = {"mean": [], "min": [], "max": []} + for k in metrics: + # valid_tokens is folded into the mean stack β€” it SUM-reduces too. + by_op.setdefault(metric_ops.get(k, "mean"), []).append(k) + + reduced: Dict[str, torch.Tensor] = {} + for op_name, keys in by_op.items(): + if not keys: + continue + stacked = torch.stack([torch.as_tensor(metrics[k], dtype=torch.float64, device=device) for k in keys]) + reduce_op = {"min": dist.ReduceOp.MIN, "max": dist.ReduceOp.MAX}.get(op_name, dist.ReduceOp.SUM) + dist.all_reduce(stacked, op=reduce_op, group=sp_group) + for i, k in enumerate(keys): + reduced[k] = stacked[i] + + metrics.update(reduced) + # valid_tokens is a Python int downstream (used as a denominator). + if "valid_tokens" in reduced: + metrics["valid_tokens"] = int(reduced["valid_tokens"].item()) return metrics @@ -724,51 +709,82 @@ 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 _accumulate_is_metrics(accumulated, new_metrics, metric_ops=None): + """Accumulate IS metrics across micro-batches. + + ``metric_ops`` tags non-mean keys; means accumulate by sum (finalized as + sum/count in _finalize_is_metrics), min/max by torch.minimum/maximum. + Values stay on-device β€” one collective + one ``.item()`` per metric in + _finalize_is_metrics. + """ if not new_metrics: return - new_metrics.pop("_n_valid_kl", None) - n_tokens = new_metrics.get("valid_tokens", 1) + metric_ops = metric_ops or {} + n_tokens = float(new_metrics.get("valid_tokens", 1)) for k, v in new_metrics.items(): - if k not in accumulated: - accumulated[k] = {"sum": 0.0, "count": 0} - if k == "valid_tokens": - accumulated[k]["sum"] += v - accumulated[k]["count"] += 1 + op = metric_ops.get(k, "mean") + v_t = torch.as_tensor(v, dtype=torch.float64, device="cuda") + if op in ("min", "max"): + entry = accumulated.get(k) + if entry is None: + accumulated[k] = {"value": v_t.clone(), "op": op} + else: + entry["value"] = ( + torch.minimum(entry["value"], v_t) if op == "min" else torch.maximum(entry["value"], v_t) + ) else: - accumulated[k]["sum"] += v * n_tokens - accumulated[k]["count"] += n_tokens + entry = accumulated.get(k) + if entry is None: + accumulated[k] = { + "sum": v_t.clone(), + "count": 1.0 if k == "valid_tokens" else n_tokens, + "op": "mean", + } + else: + entry["sum"] = entry["sum"] + v_t + entry["count"] += 1.0 if k == "valid_tokens" else n_tokens @staticmethod def _finalize_is_metrics(accumulated, result): - """All-reduce IS metrics across DP group, then add averaged values to result dict.""" + """All-reduce IS metrics across DP, then write finalized values to result. + + One SUM-allreduce for mean partial-sums concatenated with their counts; + one MIN/MAX-allreduce per non-mean group. Min/max with non-finite + reductions (every rank empty) fall back to 1.0. + """ if not accumulated: return ps = get_parallel_state() - 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") - 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") - 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 - else: - sum_t = torch.tensor(v["sum"], dtype=torch.float64, device="cuda") - count_t = torch.tensor(float(v["count"]), dtype=torch.float64, device="cuda") - 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: - result[f"is_{k}"] = v["sum"] / v["count"] + dp_group = ps.dp_group if ps.dp_enabled else None + + groups: Dict[str, list[str]] = {"mean": [], "min": [], "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="cuda") + if dp_group is not None: + sums_and_counts = torch.cat([sums, counts]) + dist.all_reduce(sums_and_counts, op=dist.ReduceOp.SUM, group=dp_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[f"is_{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 dp_group is not None: + dist.all_reduce(stacked, op=reduce_op, group=dp_group) + values = stacked.tolist() + for k, v in zip(keys, values): + result[f"is_{k}"] = v if math.isfinite(v) else 1.0 def _count_global_valid_tokens(self, micro_batches): """Count valid tokens across all micro-batches and all-reduce across DP group. @@ -789,7 +805,7 @@ def _count_active_microbatches(self, micro_batches) -> tuple[int, int]: # ========================================================================= 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. Returns (local_loss_sum, per_token_outputs_dict, is_metrics, metric_ops, model_outputs).""" params = loss_fn_params or {} return_per_token = params.get("return_per_token", True) @@ -800,8 +816,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)) + per_token_outputs = {} is_metrics = None + metric_ops = None if loss_fn in ["causallm_loss", "cross_entropy"]: labels = micro_batch.get("labels") @@ -812,8 +833,9 @@ def _compute_micro_batch_loss(self, micro_batch, loss_fn, loss_fn_params): return_per_token=return_per_token, ce_mode=self.ce_mode, lm_head_fp32=self.lm_head_fp32, + loss_reducer=token_sum_reducer, ) - 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 @@ -833,13 +855,16 @@ 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, + loss_reducer=token_sum_reducer, + 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 + 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, metric_ops, get_parallel_state().ulysses_group) # Diagnostic top-k extraction (forward-only feature, rarely used) diagnostic_topk = params.get("diagnostic_topk", 0) @@ -931,18 +956,21 @@ 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, + loss_reducer=token_sum_reducer, + 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 + 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, metric_ops, get_parallel_state().ulysses_group) 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, metric_ops, outputs # ========================================================================= # Per-sample K3 KL divergence @@ -1050,16 +1078,16 @@ def _forward_loop( # Forward pass + loss computation 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, metric_ops, outputs = self._compute_micro_batch_loss( micro_batch, loss_fn, params ) 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"local_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 + # Note: local_loss_sum is always finite even when local_valid_tokens=0, because # causallm_loss_function uses reduction="none" + manual mean with # clamp(min=1) denominator. No need to replace with zeros_like # (which would break the autograd graph and cause FSDP2 deadlocks). @@ -1105,16 +1133,12 @@ def _forward_loop( } ) - # Gradient accumulation β€” raw (unnormalized) backward. - # Normalization by total accumulated valid tokens is deferred to optim_step. - # FSDP's automatic gradient averaging is disabled (set_gradient_divide_factor(1.0) - # in torch_parallelize), so no fsdp_size compensation is needed here. - # When local_valid_tokens=0, this produces 0 gradients while preserving - # the full autograd graph through all parameters (including lm_head weight), - # which is critical for FSDP2 reduce-scatter collectives. + # Backward + reporting on the raw partial sum: cross-mb / cross-DP + # accumulation composes under SUM-allreduce, then optim_step divides + # once by global_valid_tokens. FSDP grad averaging is off + # (set_gradient_divide_factor(1.0) in torch_parallelize). 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") @@ -1124,26 +1148,21 @@ def _forward_loop( set_replay_stage("replay_backward") with self.model_bwd_context: - raw_loss.backward() + local_loss_sum.backward() - # Loss reporting (separately, no grad): compute normalized per-token loss with torch.no_grad(): - loss_report = loss.detach() * local_valid_tokens + loss_report = local_loss_sum.detach() 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) + self._accumulate_is_metrics(accumulated_is_metrics, is_metrics, 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). 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/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_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) From eee2a2dab6bdff90045bc026f9c40ac8a262c787 Mon Sep 17 00:00:00 2001 From: Ashwinee Panda Date: Sat, 9 May 2026 10:33:09 -0700 Subject: [PATCH 27/49] feat(kimi): add Kimi 2.5 support --- .codespellrc | 2 +- .gitignore | 7 + pyproject.toml | 1 + src/xorl/arguments.py | 26 +- src/xorl/checkpoint/checkpointer.py | 47 +- .../sequence_parallel/ring_attention.py | 36 +- .../distributed/sequence_parallel/strategy.py | 49 +- src/xorl/distributed/torch_parallelize.py | 67 +- src/xorl/lora/utils.py | 69 +- src/xorl/models/__init__.py | 4 + src/xorl/models/auto.py | 43 +- .../models/checkpoint_handlers/buffers.py | 17 +- src/xorl/models/layers/moe/aux_loss.py | 4 + src/xorl/models/layers/moe/routing_replay.py | 29 +- src/xorl/models/module_utils.py | 1156 ++++++++++++++++- src/xorl/models/transformers/__init__.py | 4 +- .../transformers/deepseek_v3/__init__.py | 14 + .../deepseek_v3/checkpoint_handler.py | 370 ++++++ .../deepseek_v3/configuration_deepseek_v3.py | 235 ++++ .../deepseek_v3/modeling_deepseek_v3.py | 652 ++++++++++ .../transformers/deepseek_v3/parallelize.py | 22 + .../transformers/deepseek_v3/support.py | 138 ++ .../deepseek_v3/tokenization_kimi.py | 248 ++++ .../models/transformers/qwen3_5_shared.py | 8 + src/xorl/server/runner/checkpoint/manager.py | 60 +- src/xorl/server/runner/model_runner.py | 87 +- src/xorl/server/server_arguments.py | 7 + src/xorl/trainers/model_builder.py | 29 +- src/xorl/trainers/trainer.py | 92 +- src/xorl/trainers/training_utils.py | 27 +- .../test_deepseek_v3_ep_checkpoint.py | 163 +++ .../test_deepseek_v3_checkpoint_handler.py | 272 ++++ tests/models/test_deepseek_v3_model.py | 146 +++ tests/models/test_deepseek_v3_registry.py | 125 ++ tests/models/test_kimi_tokenizer.py | 82 ++ .../models/test_lora_moe_attention_targets.py | 70 + tests/models/test_model_state.py | 33 + tests/models/test_module_utils_broadcast.py | 499 ++++++- tests/models/test_qwen3_5_apply_rotary.py | 92 ++ .../server/runner/test_checkpoint_loading.py | 87 ++ .../runner/test_deepseek_v3_lora_targets.py | 102 ++ .../test_deepseek_v3_training_guards.py | 105 ++ 42 files changed, 5162 insertions(+), 164 deletions(-) create mode 100644 src/xorl/models/transformers/deepseek_v3/__init__.py create mode 100644 src/xorl/models/transformers/deepseek_v3/checkpoint_handler.py create mode 100644 src/xorl/models/transformers/deepseek_v3/configuration_deepseek_v3.py create mode 100644 src/xorl/models/transformers/deepseek_v3/modeling_deepseek_v3.py create mode 100644 src/xorl/models/transformers/deepseek_v3/parallelize.py create mode 100644 src/xorl/models/transformers/deepseek_v3/support.py create mode 100644 src/xorl/models/transformers/deepseek_v3/tokenization_kimi.py create mode 100644 tests/distributed/test_deepseek_v3_ep_checkpoint.py create mode 100644 tests/models/test_deepseek_v3_checkpoint_handler.py create mode 100644 tests/models/test_deepseek_v3_model.py create mode 100644 tests/models/test_deepseek_v3_registry.py create mode 100644 tests/models/test_kimi_tokenizer.py create mode 100644 tests/models/test_lora_moe_attention_targets.py create mode 100644 tests/models/test_model_state.py create mode 100644 tests/models/test_qwen3_5_apply_rotary.py create mode 100644 tests/server/runner/test_checkpoint_loading.py create mode 100644 tests/server/runner/test_deepseek_v3_lora_targets.py create mode 100644 tests/trainers/test_deepseek_v3_training_guards.py 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..aaa5bd5e 100644 --- a/.gitignore +++ b/.gitignore @@ -76,3 +76,10 @@ 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/ diff --git a/pyproject.toml b/pyproject.toml index 0abc57a3..c2060434 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -18,6 +18,7 @@ dependencies = [ "packaging>=23.0,<26.0", "tenacity>=8.0.0", "torchdata>=0.8.0,<1.0", + "tiktoken", "transformers[torch]>=5.0", "psutil", "wandb", diff --git a/src/xorl/arguments.py b/src/xorl/arguments.py index e0b96e58..5cc7085e 100644 --- a/src/xorl/arguments.py +++ b/src/xorl/arguments.py @@ -455,6 +455,10 @@ 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={ @@ -814,10 +818,10 @@ 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( + load_weights_mode: Literal["broadcast", "all_ranks", "grouped", "skip"] = field( default="broadcast", 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. 'broadcast': global rank0 reads weights and broadcasts to other ranks. 'all_ranks': every rank reads weights from disk independently. 'grouped': one reader per EP-FSDP group reads and broadcasts within that group. 'skip': skip HF weight loading (use with load_checkpoint_path for DCP)." }, ) enable_full_determinism: bool = field( @@ -1102,6 +1106,14 @@ def __post_init__(self): ckpt_manager=self.ckpt_manager, ) + 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") @@ -1229,6 +1241,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 diff --git a/src/xorl/checkpoint/checkpointer.py b/src/xorl/checkpoint/checkpointer.py index c91646c6..83f69bb0 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): """ @@ -607,9 +648,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/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 1beaa791..188a8c04 100644 --- a/src/xorl/distributed/sequence_parallel/strategy.py +++ b/src/xorl/distributed/sequence_parallel/strategy.py @@ -58,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 # ------------------------------------------------------------------ # @@ -176,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 @@ -246,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) diff --git a/src/xorl/distributed/torch_parallelize.py b/src/xorl/distributed/torch_parallelize.py index 70ed5e61..cc92af42 100644 --- a/src/xorl/distributed/torch_parallelize.py +++ b/src/xorl/distributed/torch_parallelize.py @@ -19,7 +19,8 @@ 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, rank0_load_and_broadcast_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 @@ -38,6 +39,27 @@ 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: + if load_weights_mode == "skip": + logger.info_rank0("Skipping HF weight loading (weights will be loaded from DCP checkpoint).") + return + elif 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=dtensor_factory) + elif load_weights_mode == "grouped": + logger.info_rank0("Loading model weights with one reader per EP-FSDP group...") + grouped_load_weights(model, weights_path, weight_device, dtensor_factory=dtensor_factory) + else: + logger.info_rank0("Every rank reading weights from disk independently...") + all_ranks_load_weights(model, weights_path, weight_device, dtensor_factory=dtensor_factory) + + _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, @@ -417,12 +439,13 @@ def _experts_shard_placement_fn(param): 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_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 @@ -590,6 +613,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()): @@ -617,14 +641,13 @@ def build_parallelize_model( 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_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) @@ -736,6 +759,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.") @@ -761,12 +785,13 @@ def _reentrant_ckpt_with_kwargs(fn, *args, **kw): 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_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 diff --git a/src/xorl/lora/utils.py b/src/xorl/lora/utils.py index a6d55c19..a58e1b0b 100644 --- a/src/xorl/lora/utils.py +++ b/src/xorl/lora/utils.py @@ -30,6 +30,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." @@ -46,6 +76,18 @@ LORA_EXPORT_FORMATS = ("peft", "sglang_shared_outer") +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]: """ Get parent module and attribute name for a target path. @@ -167,7 +209,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) @@ -886,9 +928,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: @@ -904,13 +947,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", "qkv_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) 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 bee168cc..447f420c 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,8 @@ 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.qwen3_5.configuration_qwen3_5 import Qwen3_5Config from .transformers.qwen3_5_moe.configuration_qwen3_5_moe import Qwen3_5MoeConfig @@ -31,6 +35,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()}) @@ -55,13 +86,16 @@ def _load_local_xorl_config( 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 == "olmo2": - from .transformers.olmo2.configuration_olmo2 import Olmo2Config + from .transformers.olmo2.configuration_olmo2 import Olmo2Config # noqa: PLC0415 return Olmo2Config(**{k: v for k, v in config_dict.items() if not k.startswith("_")}) @@ -72,6 +106,9 @@ 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") @@ -153,6 +190,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'. " 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/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/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/module_utils.py b/src/xorl/models/module_utils.py index ca934df8..e97cb1d8 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,387 @@ 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("language_model."): + return key.removeprefix("language_model.") + return key + + +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 + + +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 +722,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 +874,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 +967,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 +1076,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 +1295,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 +1319,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 +1533,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 +1767,392 @@ 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) + + 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() + + _ps = get_parallel_state() + fanout_group = _get_grouped_weight_load_group(_ps) + if fanout_group is None: + logger.warning_once( + "Grouped weight loading requires EP/FSDP groups; falling back to rank0_load_and_broadcast_weights." + ) + return rank0_load_and_broadcast_weights(model, weights_path, init_device, dtensor_factory) + + 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, + ) + 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 _should_skip_dense_key(key: str) -> bool: + normalized = _normalize_checkpoint_key_for_filter(key) + return normalized is None or _is_checkpoint_expert_key(normalized) + + def _should_skip_grouped_expert_key(key: str) -> bool: + normalized = _normalize_checkpoint_key_for_filter(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_checkpoint_key_for_filter(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 +2276,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 +2423,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 +2457,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: @@ -1664,8 +2758,8 @@ 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 diff --git a/src/xorl/models/transformers/__init__.py b/src/xorl/models/transformers/__init__.py index 5af9d2d1..0edc5da7 100644 --- a/src/xorl/models/transformers/__init__.py +++ b/src/xorl/models/transformers/__init__.py @@ -1,4 +1,5 @@ from . import ( + deepseek_v3, glm4_moe, llama3, qwen2, @@ -10,11 +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/qwen3_5_shared.py b/src/xorl/models/transformers/qwen3_5_shared.py index d93a0860..a547ff19 100644 --- a/src/xorl/models/transformers/qwen3_5_shared.py +++ b/src/xorl/models/transformers/qwen3_5_shared.py @@ -51,6 +51,14 @@ def qwen3_5_apply_rotary_pos_emb( sin: torch.Tensor, interleaved: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: + 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/server/runner/checkpoint/manager.py b/src/xorl/server/runner/checkpoint/manager.py index ee05a664..acaabe5e 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 @@ -32,6 +33,7 @@ from xorl.lora.utils import get_lora_state_dict, save_lora_checkpoint from xorl.models import save_model_weights from xorl.utils import helper +from xorl.utils.device import get_device_type logger = logging.getLogger(__name__) @@ -107,6 +109,57 @@ 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) # ------------------------------------------------------------------ @@ -941,7 +994,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) @@ -950,12 +1003,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 6c3747c7..b4206883 100644 --- a/src/xorl/server/runner/model_runner.py +++ b/src/xorl/server/runner/model_runner.py @@ -23,7 +23,7 @@ import torch import torch.distributed as dist import torch.nn.functional as F -from transformers import AutoTokenizer +from transformers import AutoTokenizer, PretrainedConfig from xorl.checkpoint import build_checkpointer from xorl.data.constants import IGNORE_INDEX @@ -33,6 +33,7 @@ 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 ( TokenPartial, causallm_loss_function, @@ -54,9 +55,9 @@ 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, ) @@ -191,6 +192,11 @@ def __init__( self.model_config = config.get("model", {}) self.train_config = config.get("train", {}) self.lora_config = config.get("lora", {}) + 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") @@ -214,6 +220,7 @@ def __init__( # Multi-adapter support (initialized later if LoRA is enabled) self._adapter_manager: Optional[LoRAAdapterManager] = 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 @@ -241,6 +248,8 @@ def __init__( self._initialize_model() self._initialize_optimizer() self._initialize_checkpointer() + self._checkpoint_mgr = self._build_checkpoint_manager() + self._load_initial_checkpoint() self._initialize_contexts() # Initialize multi-adapter manager if LoRA is enabled @@ -259,6 +268,7 @@ def __init__( initialize_fresh=False, # Use current weights as the default adapter ) 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) @@ -271,17 +281,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) @@ -538,14 +537,36 @@ 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 @@ -612,6 +633,30 @@ def _initialize_contexts(self): self.train_config.get("activation_gpu_limit", None), ) + 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() + def register_lora_adapter(self, model_id: str, lr: float) -> Dict[str, Any]: """ Register a new LoRA adapter for a training run. @@ -1241,7 +1286,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] diff --git a/src/xorl/server/server_arguments.py b/src/xorl/server/server_arguments.py index a3df5ab8..e72fa973 100644 --- a/src/xorl/server/server_arguments.py +++ b/src/xorl/server/server_arguments.py @@ -564,6 +564,13 @@ def __post_init__(self): # the launcher can still use them via engine_connect_host + worker_bind_port pass + 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. diff --git a/src/xorl/trainers/model_builder.py b/src/xorl/trainers/model_builder.py index 9c9ffd84..6459f38e 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, @@ -195,6 +200,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") # ------------------------------------------------------------------ @@ -253,6 +264,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, @@ -265,7 +277,7 @@ 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( @@ -313,11 +325,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") @@ -442,8 +455,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, diff --git a/src/xorl/trainers/trainer.py b/src/xorl/trainers/trainer.py index bd4b5dcf..4242d22b 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, @@ -60,10 +64,10 @@ 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, ) @@ -74,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]: @@ -86,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 # --------------------------------------------------------------------------- @@ -165,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}") @@ -281,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] @@ -291,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( @@ -414,6 +449,12 @@ def _build_model(self) -> None: init_device=args.train.init_device, ) self.model_config = self.model.config + 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 @@ -495,7 +536,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( @@ -506,7 +547,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, @@ -573,6 +614,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 @@ -676,7 +727,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", {}) @@ -720,7 +791,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] @@ -1176,6 +1247,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 50880a13..dd6591b4 100644 --- a/src/xorl/trainers/training_utils.py +++ b/src/xorl/trainers/training_utils.py @@ -270,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(), @@ -286,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/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/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_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_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..ef7e728f --- /dev/null +++ b/tests/models/test_model_state.py @@ -0,0 +1,33 @@ +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 diff --git a/tests/models/test_module_utils_broadcast.py b/tests/models/test_module_utils_broadcast.py index bb4681c5..5c03de39 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,227 @@ 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"] + + +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_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_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)), + ) + + module_utils.grouped_load_weights(_DummyModel(), "dummy-weights", init_device="cpu") + + assert len(called) == 1 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..e476c814 --- /dev/null +++ b/tests/models/test_qwen3_5_apply_rotary.py @@ -0,0 +1,92 @@ +import pytest +import torch + +from xorl.models.layers.rope import rotate_half +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_interleaved( + q: torch.Tensor, k: torch.Tensor, cos_halved: torch.Tensor, sin_halved: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor]: + """HF-style reference for interleaved-layout q/k with halved cos/sin. + + 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_interleaved_matches_hf_reference(): + 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_interleaved(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) + + +def test_non_interleaved_unchanged(): + """Non-interleaved path must be unaffected by the fix.""" + 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_out, k_out = qwen3_5_apply_rotary_pos_emb(q, k, cos, sin, interleaved=False) + + cos_u = cos.unsqueeze(2) + sin_u = sin.unsqueeze(2) + expected_q = q * cos_u + rotate_half(q) * sin_u + expected_k = k * cos_u + rotate_half(k) * sin_u + torch.testing.assert_close(q_out, expected_q, atol=1e-6, rtol=1e-6) + torch.testing.assert_close(k_out, expected_k, atol=1e-6, rtol=1e-6) 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_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/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()) From 5d9ce37720907758985c3e18eb16a42121f40128 Mon Sep 17 00:00:00 2001 From: Qingyang Wu Date: Tue, 12 May 2026 15:39:25 -0700 Subject: [PATCH 28/49] fix(tests): align Qwen3.5-MoE layer_types fixture with num_hidden_layers The fixture passed a 2-element `layer_types` pattern alongside `num_hidden_layers=40`. Upstream transformers' `Qwen3MoeConfig` now runs a strict `validate_layer_type` validator (via huggingface_hub dataclasses) that requires `len(layer_types) == num_hidden_layers`, which rejected the short pattern and failed gpu-tests on main. Build the expected per-layer pattern from `full_attention_interval` and assert against the full list. No production code change needed -- real HF configs already provide a full-length `layer_types`. --- tests/models/test_qwen3_5_registry.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/tests/models/test_qwen3_5_registry.py b/tests/models/test_qwen3_5_registry.py index bec9f4e4..5fe1bf4e 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 From 69c84242cc4c571c52638791c79b9c361bf8fd22 Mon Sep 17 00:00:00 2001 From: Qingyang Wu Date: Thu, 14 May 2026 10:27:53 -0700 Subject: [PATCH 29/49] refactor: tighten broad except blocks; preserve causes First batch toward. - distributed/fsdp2/clip_grad_norm.py: narrow `except Exception` around dist.get_world_size to RuntimeError (the only expected failure when the process group hasn't been initialized). - arguments.py: narrow the git-detection except to (subprocess.SubprocessError, FileNotFoundError, OSError); narrow the two get_type_hints sites to (NameError, TypeError, AttributeError) and chain the cause via `from e` so the underlying type-resolution failure isn't lost. - data/prepare/shared.py:515: keep the broad catch (datasets raises a zoo of types) but include type and message in the log so silent "not found" failures are debuggable. Other broad-except sites tracked under. --- src/xorl/arguments.py | 13 ++++++++----- src/xorl/data/prepare/shared.py | 8 ++++++-- src/xorl/distributed/fsdp2/clip_grad_norm.py | 5 ++++- 3 files changed, 18 insertions(+), 8 deletions(-) diff --git a/src/xorl/arguments.py b/src/xorl/arguments.py index 5cc7085e..14cbf65a 100644 --- a/src/xorl/arguments.py +++ b/src/xorl/arguments.py @@ -47,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 @@ -1373,8 +1376,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: @@ -1410,8 +1413,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/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/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)) From 03658d60ee70b07be3ae6142ddafa73057933654 Mon Sep 17 00:00:00 2001 From: Qingyang Wu Date: Thu, 14 May 2026 10:28:15 -0700 Subject: [PATCH 30/49] refactor: use sys.stderr.write for CLI startup errors instead of print MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two CLI entrypoints (xorl.cli.preprocess and the runner dispatcher in server/runner/setup) print error messages to stdout before falling through to sys.exit(1). Errors belong on stderr β€” switching to sys.stderr.write keeps the user-visible behavior while letting us drop the per-file T201 ruff ignores PR added (in a follow-up after merges). Refs. --- src/xorl/cli/preprocess.py | 4 ++-- src/xorl/server/runner/setup.py | 6 ++++-- 2 files changed, 6 insertions(+), 4 deletions(-) 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/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 From 2b43ffba7b183571b3d107f7af4666a80f125084 Mon Sep 17 00:00:00 2001 From: AMAN SINGHAL <65446479+AmanSinghal927@users.noreply.github.com> Date: Thu, 14 May 2026 13:56:05 -0400 Subject: [PATCH 31/49] Aman/gpt oss ep support MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Add GPT-OSS model support with Expert Parallelism Adds full GPT-OSS (20b/120b) model implementation with EP support: Model implementation: - GptOssForCausalLM with custom clamped SwiGLU activation, attention sinks, alternating sliding window, YaRN RoPE, and per-expert biases - Checkpoint handler supporting both original (MXFP4) and HF (BF16) checkpoint formats with weight deinterleaving - GptOssConfig with from_hf_config() supporting both 20b and 120b Generic MoE infrastructure changes (model-agnostic): - Added optional gate_up_bias, down_bias, and act_fn parameters to native backend's _run_experts_grouped_mm and EP compute functions - Threaded bias/act_fn through MoEExperts.forward() and _ep_forward() via getattr with None defaults (zero impact on existing models) - Fixed _native_ep_fused wrapper to pass through expert_scores and **kwargs (pre-existing bug where expert_scores was silently dropped) Training configs and test script: - Training configs for 20b (EP=8) and 120b (EP=8) with Muon optimizer - EP correctness test comparing logprobs across EP configurations * add ep support changes * apply ruff format and drop unused import to fix lint CI * fix lint CI: apply ruff --fix-only and format - remove unused os import in gpt_oss_gsm8k_hf_logprobs.py - sort imports in gpt_oss_gsm8k_xorl_logprobs.py and test_gpt_oss_ep_correctness.py - drop f-string prefix on strings without placeholders - reformat moe/backend/__init__.py long signatures * mark gpt-oss scripts as executable to match shebangs * feat(gpt-oss): apply sinks under flash_attention_3 backend The eager attention path applied the learned per-head sink; the flash_attention_3 path silently dropped it, producing mathematically incorrect GPT-OSS outputs whenever a user configured FA3. Add a GPT-OSS-local autograd wrapper that fuses the sink into FA3 via sigmoid(lse - sink) (equivalent to eager's concat-softmax-drop), and route flash_attention_3 through it. Supports both batched and packed-varlen inputs; FA4 and other non-eager backends now raise NotImplementedError instead of dropping sinks silently. Verified against eager reference: forward and all four gradients (dq, dk, dv, dsink) match at bf16 noise floor (cos >= 0.99992) across causal, GQA, sliding-window, head_dim=128, and multi-sequence packed cases. * refactor(moe): fold GPT-OSS clamped SwiGLU into hidden_act dispatch PR introduced a parallel act_fn_override callable escape hatch on MoEExperts to plumb GPT-OSS's clamped SwiGLU through the MoE backends, bypassing the hidden_act string dispatch established by. This left two coexisting activation mechanisms and forced torch.compile to deal with a Python callable closure argument β€” graph-break risk. Extend instead: register "clamped_swiglu" as a recognized hidden_act kind in normalize_hidden_act and SUPPORTED_HIDDEN_ACTS, with a _clamped_swiglu_act helper in native.py (alpha=1.702, limit=7.0 hardcoded per GPT-OSS spec). Drop act_fn / act_fn_split / act_fn_override plumbing from experts.py, eager.py, native.py, backend/__init__.py, and the GPT-OSS model. GptOssMoEBlock sets self.experts.hidden_act = "clamped_swiglu" β€” same dispatch as every other model, fully compile-friendly. Eager-path pre/post snapshot bitwise-identical for out, dx, and all parameter grads on a GPT-OSS MoE block with non-zero gate_up_bias / down_bias. Native-path forward+backward also verified finite. Quack / triton EP adapters still raise NotImplementedError loudly for biases or clamped_swiglu β€” they don't implement those kernels. * refactor(gpt-oss): replace GptOssRMSNorm with shared RMSNorm GptOssRMSNorm existed to do the weight multiply in fp32 before casting to bf16 β€” a precision-matching copy of HF's reference. It worked but (1) ignored xorl's rmsnorm_mode CLI flag entirely, so enabling compile mode was silently inert on every norm in the model; (2) left a separate eager-only class drifting from the shared RMSNorm abstraction. Drop the custom class and use the shared RMSNorm (default "native" mode, F.rms_norm under the hood; honours set_rmsnorm_mode for eager/compile). Snapshot-verified on a full GPT-OSS decoder layer (eager attn + eager MoE + non-zero biases / sinks): the forward output, input gradient, and 13/14 parameter gradients are bitwise-identical. Only input_layernorm. weight grad differs β€” max_abs=6e-5, rel=2e-5 β€” a precision-level artefact of F.rms_norm's fused weight grad path, orders of magnitude below the bf16 training noise floor (~3e-3 relative). Net: βˆ’40 lines, rmsnorm_mode="compile" now meaningful for GPT-OSS. * refactor(moe): extract hidden_act dispatch to a dedicated module Activation logic was scattered: SUPPORTED_HIDDEN_ACTS + normalize_hidden_act + check_hidden_act_supported lived in ops/moe/triton.py (next to triton kernel internals), per-activation helpers (_default_swiglu_act, _clamped_swiglu_act, _CLAMPED_SWIGLU_*) lived in the native backend, and eager.py had to reach across to native for the clamped variant. Move everything to src/xorl/ops/moe/activations.py. It owns: - Per-activation implementations (silu_swiglu, gelu_tanh_glu, clamped_swiglu) and their constants (CLAMPED_SWIGLU_ALPHA, CLAMPED_SWIGLU_LIMIT). - MOE_ACTIVATIONS registry + SUPPORTED_HIDDEN_ACTS derived from its keys. - normalize_hidden_act / check_hidden_act_supported. - apply_moe_activation(hidden_act, gate, up) β€” the single dispatch helper, written as explicit if-chain so torch.compile specializes on the string. Native + eager backends now just call apply_moe_activation; no more inline if/elif chains or cross-backend imports. Triton + quack keep re-exporting the validation helpers from their old import paths for backward compat. Adding a new activation is now a three-step change in one file: write the function, add it to MOE_ACTIVATIONS, add the string to normalize_hidden_act. Backends opt in by listing the name in their per-class SUPPORTED_HIDDEN_ACTS. Eager-path snapshot on a full GPT-OSS decoder layer bitwise-identical to pre-refactor (13/13 tensors β€” forward, input grad, every parameter grad). * Fix GPT-OSS reload and attention validation * adding tests * apply ruff format to fix lint CI * mark gpt-oss scripts as executable (reverted by merge) * fix tests broken by qingyang/gpt-oss-supplement merge - test_adapter_source_forwards_expert_scores: collapse whitespace before the string match so ruff-format rewrapping the _native_ep_compute call across lines no longer breaks the assertion. - test_build_foundation_model_rejects_non_eager_gpt_oss_attention: the merge added flash_attention_3 sink support, so 'flash_attention_3' is now accepted. Switch the rejected case to 'native' and rename to test_build_foundation_model_rejects_unsupported_gpt_oss_attention. * restore executable bit on gpt-oss scripts * Add softmax auxiliary (Z-)loss on LM-head logits * Add DeepSeek V4 distributed Muon paths * Add DeepSeek V4 distributed Muon paths Implements two opt-in features from DeepSeek V4 Β§3.5.1, both default off: 1. Full-gradient Muon NS (muon_distributed_mode=full_gradient) - Replaces shard-local NS on FSDP2/EP DTensor params with NS on the all-gathered full matrix, recovering the exact Muon update direction. - Momentum/Nesterov stay on the local shard (linear in grad, commutes with sharding) so the optimizer-state buffer stays at local-shard size. - LR adjustment uses the global matrix shape (fixes a bug where adjust_lr_fn saw the local-shard shape under shard_local mode). - Knapsack matrix-to-rank assignment and bounded ZeRO width are not in this PR; every rank in the param's mesh runs NS redundantly. 2. BF16 stochastic-rounded a2a + FP32 local sum (moe_grad_reduce_mode =bf16_a2a_fp32_sum) - Custom FSDPModule.set_custom_reduce_scatter on EP-experts modules: stochastic-round FP32β†’BF16, dist.all_to_all_single, sum locally in FP32. Halves comm volume vs FP32 reduce-scatter while preserving FP32 accumulation precision. - Stochastic rounding (xorl.optim.stochastic_round) operates on the FP32 bit pattern; unbiased in expectation. Tests: - CPU: stochastic round dtype, unbiasedness, bracket bounds, determinism. - Distributed (2 GPU): full-gradient Muon matches a single-rank oracle; shard-local provably differs on multi-rank. - Distributed (4 GPU): BF16 a2a reduce-scatter matches FP32 reduce-scatter within BF16 ulp bounds, with unbiased mean over many trials. - Distributed (2 GPU): end-to-end FSDP2 hook smoke test. * Address self-review: BF16 hook preconditions, plan typing, 3D EP test - Assert mp_policy.reduce_dtype=fp32 and gradient_divide_factor=1.0 on every expert FSDPModule before installing BF16StochasticAllToAllReduceScatter. Both are correctness preconditions: a non-fp32 reduce_dtype would surface a runtime error during the first backward (far from configuration), and a factor != 1.0 means FSDP applies predivide before the hook and skips the postdivide because the reduce is custom β€” silently under-weighting grads. Document the same preconditions on the hook's docstring. - Type _MuonUpdatePlan.placements as Optional[Tuple[Placement,...]] and device_mesh as Optional[DeviceMesh] via TYPE_CHECKING imports, replacing the loose Optional[object] / Optional[Tuple] placeholders. - Add a 3D-EP-experts (Shard(1) on the dp mesh) variant to the existing full_gradient distributed test, exercising the deferred-reshape branch in _muon_step that the previous tests only covered for 2D Shard(0). Both full_gradient (matches single-rank oracle) and shard_local (provably differs) layouts are checked. * Apply ruff-format to test files post-merge * fix(lora): merge in fp32 and cast the sum once * LoRA: merge in fp32 and cast the sum once Current ``LoraLinear.merge_weights`` and ``MoEExpertsLoRA.merge_weights`` round the LoRA delta per element to the weight dtype before adding: W.add_(delta.to(W.dtype)) This loses ~1 weight-dtype ULP per element of Ξ”. On bf16 weights with MoE top-k routing downstream, tiny per-element drift can flip expert selection at the router (argmax top-k is discontinuous), which cascades through the stack. Empirically, on Qwen3-30B-A3B with random lora_B ~N(0, 0.005) the naive variant degrades K3 by ~20x vs the fp32-sum-then-cast variant: std=0.005 naive K3=1.9e-1 fp32-cast-once K3=1.0e-2 (18x better) The change is a one-liner: upcast W to fp32, add the fp32 Ξ”, cast the sum back once. W.data.copy_((W.to(fp32) + delta).to(W.dtype)) Same memory as before (result lands in W.dtype). Strictly β‰₯ precision of the naive variant β€” bit-exact in fp32, more faithful in bf16/fp16. On realistic trained LoRAs (password-memorization adapters, std << 0.005) both variants produce identical greedy outputs; this change widens the "safe" margin and is the right default. New tests (tests/models/test_lora_merge_fp32_cast_once.py): - zero-LoRA merge is bit-exact - merged weight matches the fp32-sum-then-cast reference bit-for-bit - fp32-cast-once error ≀ naive error (vs true fp32 merged value) - same invariants on MoEExpertsLoRA (fused gate_up + down) All existing LoRA tests still pass (54 total). * style: apply ruff fixes to lora merge test --------- Co-authored-by: Ashwinee Panda * fix: Muon/MoE backward perf regression on Qwen3.5-style MoE * Fix MoE EP backward perf regression for train_router=False Two related changes that recover lost throughput on Qwen3.5-style MoE training when train_router=False (the default for ep_dispatch='alltoall' and the only supported setting for ep_dispatch='deepep'): 1. Triton/Quack EP group GEMM backward (src/xorl/ops/moe/triton.py:127, src/xorl/ops/moe/quack.py:316): skip the extra full down-GEMM that computes grad_expert_scores when expert_scores does not require a gradient. With train_router=False MoEBlock detaches routing_weights upstream so ctx.needs_input_grad[5] is False, making the GEMM purely wasted work in backward. train_router=True still computes grads normally. 2. Make routing replay's record/pop of routing_weights opt-in via a new model arg record_routing_weights (default True for safety). When attention forward is deterministic across checkpoint recompute, the regathered routing_weights match the recorded ones, so the cache is unnecessary and disabling it avoids per-MoE-layer pinned CPU allocations + D2H/H2D copies on every step. The arg is threaded through ModelArguments / ServerArguments -> Trainer / ModelRunner -> build_training_model -> build_foundation_model -> config.record_routing_weights -> MoEBlock.from_config. * Switch Muon ns_algorithm default to gram_newton_schulz + batch standard NS The standard_newton_schulz path was the dominant cause of a ~2x training-step regression on Qwen3.5-35B-A3B vs the prior baseline. With grad shape [E, H, I] (E=local experts) it ran an O(E)-deep Python list comprehension calling the upstream 2D _zeropower_via_newtonschulz once per matrix, emitting ~14k extra kernel launches per optimizer step (40 MoE layers x 3 weight pieces x 8 experts x 5 NS steps x ~3 matmuls each). Two complementary fixes: 1. Default the Muon Newton-Schulz backend to gram_newton_schulz in arguments / server_arguments / optimizer factory / Muon.__init__. gram_newton_schulz already batches across experts via baddbmm and recovered the full ~2x throughput in a 1-pod test (tok/s 12k -> 27k, matching the prior baseline's 27,651 tok/s). 2. Add _batched_zeropower_via_newtonschulz that batches the upstream standard NS recurrence with bmm/baddbmm for users who pin ns_algorithm='standard_newton_schulz'. Bit-exact vs the per-matrix loop on representative shapes; ~2.76x faster on a single CUDA microbench at [8, 2048, 768]. The existing test_muon_standard_newton_schulz_preserves_batched_leading_dims test was monkeypatching the per-matrix _zeropower_via_newtonschulz; updated it to monkeypatch the batched variant and check the [B, H, I] flattening contract end-to-end. (cherry picked from commit 130620902de06418c2d180f635338a1143d0e05a) * ci: re-trigger workflows after title update --------- Co-authored-by: Qingyang Wu * Add Cautious Weight Decay (CWD) * Add Cautious Weight Decay (CWD) Implement Cautious Weight Decay from Chen et al. (arXiv:2510.12402): mask the decoupled weight-decay term by I(u_t * x_t >= 0) so decay only acts on coordinates whose optimizer update aligns with the parameter. The original objective is preserved (no implicit regularizer), and the modification is a one-line change with no extra hyperparameters. Plumbed via a top-level `cautious_weight_decay` flag on TrainArgs and ServerTrainArgs. Supported for anyprecision_adamw, signsgd, muon, and adamw (which auto-routes to AnyPrecisionAdamW with fp32 state since the fused torch.optim.AdamW kernel has no per-coordinate decay hook). SGD explicitly rejects the flag. Mask sign proxy chosen per optimizer: - AdamW family: exp_avg (denominator is positive, so sign matches u_t) - SignSGD: grad - Muon: post-Newton-Schulz update tensor (the actual u_t) Tests cover the mask helper, each optimizer's masked path, the cautious==standard equivalence when all signs align, the build_optimizer routing for adamw/signsgd/anyprecision_adamw/muon/sgd, and the post-NS-update mask for Muon. * Address review: rename Muon test to assert hand-computed reference * ci: replace custom CalVer workflow with semantic-release Drops the hand-rolled tag/changelog/release bash in version-bump.yml and delegates to semantic-release. Commit types now drive the version bump (feat β†’ minor, fix/perf/revert β†’ patch, BREAKING CHANGE or `!` β†’ major), which is the point of using Conventional Commits in the first place. The existing CalVer tags (latest: v26.4.20) are valid SemVer and are treated as the baseline, so the next release derives from there. Plugins are limited to commit-analyzer, release-notes-generator, and @semantic-release/github β€” no git/changelog plugins, so nothing is pushed back to main (the branch ruleset blocks that anyway). Release notes live on the GitHub Release, matching the prior behavior. * fix: resolve dp_sp alias in loss_group/loss_mesh Co-authored-by: Qingyang Wu * chore(lora): remove unused stacked LoRA helpers, drop lora_utils.py * chore(lora): remove unused merge_lora_weights_stacked / unmerge_lora_weights_stacked These two helpers in src/xorl/ops/group_gemm/kernel/lora_utils.py have no in-tree callers and are only re-exported through xorl.lora and xorl.ops.group_gemm.kernel. They date back to the original MoE+LoRA+EP support commit (74be866) and were superseded by MoEExpertsLoRA.merge_weights (the path PR actually patched for fp32-cast-once precision). Removes the function definitions and prunes them from both __init__.py re-export lists. get_lora_delta_weight_stacked is also unused but is left in place for now; flag separately if it should go too. * chore(lora): remove unused stacked LoRA helpers Removes four dead helpers from src/xorl/ops/group_gemm/kernel/lora_utils.py: - merge_lora_weights_stacked - unmerge_lora_weights_stacked - init_lora_weights_stacked - get_lora_delta_weight_stacked None have any in-tree callers (verified by grep across.py /.yaml /.yml / .toml /.json /.md plus check for "import *" forms). They've been dead since the original MoE+LoRA+EP commit (74be866) and were superseded by MoEExpertsLoRA.merge_weights / per-module LoRA paths. The remaining helper, compute_lora_scaling, is kept β€” it is used by src/xorl/models/layers/moe/lora.py and src/xorl/qlora/modules/moe_experts.py. Both __init__.py re-export lists are pruned accordingly. Risk: external consumers (e.g. SGLang LoRA export, internal tooling) that import these symbols from xorl.lora or xorl.ops.group_gemm.kernel will break their import. Worth a sibling-repo grep before merging. * chore(lora): remove unused stacked LoRA helpers, fold compute_lora_scaling into kernel __init__ Removes four dead helpers from src/xorl/ops/group_gemm/kernel/lora_utils.py: - merge_lora_weights_stacked - unmerge_lora_weights_stacked - init_lora_weights_stacked - get_lora_delta_weight_stacked Verified by repo-wide grep across.py /.ipynb /.yaml /.toml /.json / .md /.sh /.txt plus dynamic-ref and `import *` checks: none of the four have any in-tree caller. They've been dead since the original MoE+LoRA+EP commit (74be866); the live MoE-LoRA path is MoEExpertsLoRA.merge_weights, which is the path actually patched. The only remaining helper, compute_lora_scaling, is folded directly into xorl/ops/group_gemm/kernel/__init__.py and lora_utils.py is deleted β€” keeping a one-function file for a 5-line scaling helper isn't worth it. Updates the one direct submodule importer (xorl.qlora.modules.moe_experts) to import from the package instead of the now-deleted submodule. The other caller (xorl.models.layers.moe.lora) already imported from the package and needs no change. Risk: external consumers (SGLang LoRA export path, internal tooling) that import any of the four removed symbols from xorl.lora or xorl.ops.group_gemm.kernel will break their import. Cross-repo grep recommended before merging. --------- Co-authored-by: Ashwinee Panda * Fix nccl_broadcast TCPStore master deadlock Master TCPStore was created with the PyTorch default wait_for_workers=True, which blocks the constructor until world_size-1 workers connect. The design in init_nccl_group is to create a non-blocking listening master first, then fire /init_weights_update_group at the inference endpoints (workers) from a background thread, then complete the NCCL rendezvous in the main thread via _new_process_group_helper. With the default, the master constructor blocks before the background thread can even start, so workers never receive the HTTP request and sync_inference_weights deadlocks until the request times out. Pass wait_for_workers=False so the master returns as soon as it is listening; the actual rendezvous still synchronizes via _init_training_process_group. Update the unit-test fake TCPStore to accept extra kwargs. * bug fixes * bug fix in tests * retrigger CI * gitignore comment * remove dead code * remove dead code * assertions for moe act = True in native backend * standardizing native implementation for local and ep clamped swiglu * fix backend registry bugs for moe * fix bugs in expert.py * cleanup modeling_gpt_oss.py * fix: apply pre-commit formatting * chore: remove gpt-oss debugging scripts from repo * chore: apply ruff import-order fix * chore: remove gpt-oss support test from repo * test script for comparing eager and fa3 slliding window and vanilla attnetion with sinks * test(moe): replace source-grep EP test with inspect.signature contract * linter check * move compare_gpt_oss_attn_block to gpt_oss subdir * remove unused Qwen2Config top-level import in auto.py --------- Co-authored-by: Qingyang Wu Co-authored-by: Ashwinee Panda Co-authored-by: Conner Manuel <57027354+connermanuel@users.noreply.github.com> --- .../dummy/configs/full/gpt_oss_120b_ep8.yaml | 54 ++ .../dummy/configs/full/gpt_oss_20b_ep8.yaml | 55 ++ .../configs/full/gpt_oss_120b_full.yaml | 45 + src/xorl/models/auto.py | 24 + .../models/layers/moe/backend/__init__.py | 51 +- src/xorl/models/layers/moe/backend/eager.py | 22 +- src/xorl/models/layers/moe/backend/native.py | 89 +- src/xorl/models/layers/moe/experts.py | 18 +- .../models/transformers/gpt_oss/__init__.py | 8 + .../gpt_oss/checkpoint_handler.py | 410 +++++++++ .../gpt_oss/configuration_gpt_oss.py | 204 +++++ .../gpt_oss/flash_sink_attention.py | 315 +++++++ .../transformers/gpt_oss/modeling_gpt_oss.py | 858 ++++++++++++++++++ .../transformers/gpt_oss/parallelize.py | 52 ++ src/xorl/ops/moe/activations.py | 107 +++ src/xorl/ops/moe/quack.py | 3 +- src/xorl/ops/moe/triton.py | 29 +- src/xorl/trainers/trainer.py | 3 + tests/ops/test_ep_adapter_wrappers.py | 45 +- tests/ops/test_moe_gkn_format.py | 2 +- 20 files changed, 2331 insertions(+), 63 deletions(-) create mode 100644 examples/local/dummy/configs/full/gpt_oss_120b_ep8.yaml create mode 100644 examples/local/dummy/configs/full/gpt_oss_20b_ep8.yaml create mode 100644 examples/server/configs/full/gpt_oss_120b_full.yaml create mode 100644 src/xorl/models/transformers/gpt_oss/__init__.py create mode 100644 src/xorl/models/transformers/gpt_oss/checkpoint_handler.py create mode 100644 src/xorl/models/transformers/gpt_oss/configuration_gpt_oss.py create mode 100644 src/xorl/models/transformers/gpt_oss/flash_sink_attention.py create mode 100644 src/xorl/models/transformers/gpt_oss/modeling_gpt_oss.py create mode 100644 src/xorl/models/transformers/gpt_oss/parallelize.py create mode 100644 src/xorl/ops/moe/activations.py 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/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/src/xorl/models/auto.py b/src/xorl/models/auto.py index 447f420c..f2a9f919 100644 --- a/src/xorl/models/auto.py +++ b/src/xorl/models/auto.py @@ -21,6 +21,7 @@ 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 ( @@ -94,6 +95,8 @@ def _load_local_xorl_config( 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 @@ -102,6 +105,19 @@ def _load_local_xorl_config( 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. @@ -230,6 +246,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/layers/moe/backend/__init__.py b/src/xorl/models/layers/moe/backend/__init__.py index 34d929d7..51dc7ba9 100644 --- a/src/xorl/models/layers/moe/backend/__init__.py +++ b/src/xorl/models/layers/moe/backend/__init__.py @@ -53,8 +53,20 @@ from xorl.ops.moe.triton import TritonEPGroupGemm def _triton_ep_apply( - permute_tokens, cumsum, gate_up_proj, down_proj, intermediate_size, expert_scores=None, hidden_act="silu" + 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 ) @@ -63,13 +75,25 @@ def _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, hidden_act="silu" + 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 ) @@ -78,15 +102,30 @@ def _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, hidden_act="silu" + 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) + 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 6ea7a361..3ad399ab 100644 --- a/src/xorl/models/layers/moe/backend/native.py +++ b/src/xorl/models/layers/moe/backend/native.py @@ -130,6 +130,8 @@ def _run_experts_grouped_mm( 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``. @@ -137,6 +139,10 @@ def _run_experts_grouped_mm( 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. + + 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 @@ -156,22 +162,42 @@ def _run_experts_grouped_mm( 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) - - # GLU: act(gate) * up - h = gate_out * up_out + h = gate_out * up_out # down: h @ down_proj -> (tokens, hidden) - # expert_scores applied AFTER down GEMM (not before) for bf16 consistency out = torch._grouped_mm( h, down_proj.to(compute_dtype), 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) @@ -215,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, @@ -225,12 +260,17 @@ def _native_expert_forward_impl( 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 @@ -271,7 +311,13 @@ 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) β€” NO expert_scores inside GEMM + # 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, up_proj, @@ -280,6 +326,8 @@ def _native_expert_forward_impl( padded_counts, 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 @@ -302,7 +350,7 @@ def _native_expert_forward_impl( # --------------------------------------------------------------------------- -SUPPORTED_HIDDEN_ACTS = frozenset({"silu", "gelu_tanh"}) +SUPPORTED_HIDDEN_ACTS = frozenset({"silu", "gelu_tanh", "clamped_swiglu"}) def _select_compiled_fn(hidden_act: str = "silu"): @@ -311,6 +359,8 @@ def _select_compiled_fn(hidden_act: str = "silu"): # 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 @@ -326,7 +376,11 @@ def native_expert_forward( 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) @@ -340,6 +394,8 @@ def native_expert_forward( num_experts, _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"), ) @@ -547,10 +603,15 @@ def native_ep_compute( 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`` with fused gate+up GEMM. Same interface as ``TritonEPGroupGemm.apply()`` and ``QuackEPGroupGemm.apply()``. + + 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 @@ -570,10 +631,22 @@ def native_ep_compute( counts, padded_counts, num_local_experts, permute_tokens.shape[0], padded_tokens.device ) ] = expert_scores.to(padded_tokens.dtype) + + 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 + 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, ) return _unpad(out_padded, counts, padded_counts, num_local_experts, permute_tokens.shape[0]) diff --git a/src/xorl/models/layers/moe/experts.py b/src/xorl/models/layers/moe/experts.py index 94af2a93..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. @@ -76,6 +79,11 @@ def __init__( 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" self.deepep_buffer_size_gb: float = 2.0 @@ -120,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 @@ -146,6 +156,8 @@ def forward( 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 @@ -242,6 +254,8 @@ def _ep_forward( 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 @@ -276,6 +290,8 @@ def _ep_forward_debug(self, dispatch_fn, combine_fn, compute_fn, dispatch_kwargs 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/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/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 f0ed194e..27742e33 100644 --- a/src/xorl/ops/moe/quack.py +++ b/src/xorl/ops/moe/quack.py @@ -6,7 +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.triton import _moe_gate_activation, _moe_gate_activation_backward, check_hidden_act_supported +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: diff --git a/src/xorl/ops/moe/triton.py b/src/xorl/ops/moe/triton.py index e20cadf9..8a7ff5ce 100644 --- a/src/xorl/ops/moe/triton.py +++ b/src/xorl/ops/moe/triton.py @@ -1,6 +1,13 @@ import torch 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 @@ -13,28 +20,6 @@ ) -# Canonical activation kinds understood by MoE ops. Upstream `hidden_act` -# strings (e.g. "gelu_pytorch_tanh") are normalized to one of these. -SUPPORTED_HIDDEN_ACTS: frozenset[str] = frozenset({"silu", "gelu_tanh"}) - - -def normalize_hidden_act(hidden_act: str | None) -> str: - """Normalize a HF-style ``hidden_act`` string to an 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" - 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 _moe_gate_activation(gate_output: torch.Tensor, hidden_act: str = "silu") -> torch.Tensor: """Apply gate activation by kind.""" if hidden_act == "gelu_tanh": diff --git a/src/xorl/trainers/trainer.py b/src/xorl/trainers/trainer.py index 4242d22b..7f40cb37 100644 --- a/src/xorl/trainers/trainer.py +++ b/src/xorl/trainers/trainer.py @@ -449,6 +449,9 @@ def _build_model(self) -> None: 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, diff --git a/tests/ops/test_ep_adapter_wrappers.py b/tests/ops/test_ep_adapter_wrappers.py index b9dd5f04..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 @@ -165,24 +167,31 @@ def test_adapter_forwards_expert_scores(monkeypatch, backend_type, class_name): 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. +# --------------------------------------------------------------------------- + +# Required explicit parameters for every EP_EXPERT_COMPUTE entry. +_REQUIRED_EP_PARAMS = ("expert_scores", "hidden_act") - Parses the source to confirm the apply() / compute() calls include expert_scores. - This catches silent-drop bugs without needing to load the full module. - """ - source = _BACKEND_INIT_PATH.read_text() - assert "expert_scores" in source and "_QuackEPGroupGemm.apply(" in source, ( - "_QuackEPGroupGemm.apply() call missing from backend/__init__.py" +@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}" ) - # Quack EP shim must forward expert_scores (and hidden_act) through. - assert ( - "_QuackEPGroupGemm.apply(\n permute_tokens, cumsum, gate_up_proj, down_proj, intermediate_size, expert_scores, hidden_act\n )" - in source - ), "_quack_ep_fused does not forward expert_scores to _QuackEPGroupGemm.apply()" - # Native EP shim must forward expert_scores (and hidden_act) through. - assert ( - "_native_ep_compute(permute_tokens, cumsum, gate_up_proj, down_proj, expert_scores, hidden_act=hidden_act)" - in source - ), "_native_ep_fused does not forward expert_scores to _native_ep_compute()" diff --git a/tests/ops/test_moe_gkn_format.py b/tests/ops/test_moe_gkn_format.py index c0673f96..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) From 740d7a7d0c1941bd2e6de5636716b4498777ab96 Mon Sep 17 00:00:00 2001 From: Ashwinee Panda Date: Thu, 14 May 2026 16:39:38 -0700 Subject: [PATCH 32/49] feat: add P2P (Mooncake) transport backend * feat(weight-sync): add Mooncake P2P backend * test(weight-sync): add sync stress harness * fix(weight-sync): fail closed on P2P transfer errors * fix(weight-sync): pin mooncake and harden p2p routing * fix(weight-sync): return measured transfer time * test(weight-sync): report sync timing breakdown * fix(models): expand qwen3.5 layer types * fix(weight-sync): prefer fast endpoint health probe * fix(weight-sync): stabilize async p2p medium batches * fix(weight-sync): fail closed on p2p finalization * fix(weight-sync): clean up p2p diagnostics * fix(weight-sync): speed up p2p fp8 formatting * fix(weight-sync): stream fp8 cpu workspace transfers * fix(weight-sync): honor remote sync timeout * fix(weight-sync): bound fp8 cpu workspace staging * chore(weight-sync): add sparse delta probe * fix(weight-sync): handle qwen3.5 p2p linear attention slices * fix(weight-sync): support kimi p2p receiver layouts * fix(weight-sync): avoid coalescing unregistered p2p regions * Fix P2P weight sync edge cases * fix(weight-sync): restart cached p2p prepare on fallback * fix(weight-sync): address p2p fp8 review issues * Fix P2P weight sync state handling * Fix executable bit for weight sync probe --------- Co-authored-by: Qingyang Wu --- pyproject.sglang.toml | 2 + pyproject.toml | 2 + .../qwen3_5/configuration_qwen3_5.py | 35 +- .../qwen3_5_moe/configuration_qwen3_5_moe.py | 35 +- src/xorl/server/api_server/api_types.py | 8 + .../server/api_server/inference_endpoints.py | 3 + .../server/api_server/orchestrator_client.py | 8 +- src/xorl/server/backend/remote.py | 2 +- src/xorl/server/orchestrator/orchestrator.py | 4 +- .../server/orchestrator/request_processor.py | 2 + src/xorl/server/protocol/api_orchestrator.py | 4 + src/xorl/server/server_arguments.py | 6 +- src/xorl/server/utils/zmq_channels.py | 28 +- src/xorl/server/weight_sync/README.md | 113 +- .../server/weight_sync/backends/__init__.py | 6 +- src/xorl/server/weight_sync/backends/base.py | 22 +- .../weight_sync/backends/nccl_broadcast.py | 3 +- src/xorl/server/weight_sync/backends/p2p.py | 1945 +++++++++++++++ .../server/weight_sync/endpoint_manager.py | 24 +- src/xorl/server/weight_sync/handler.py | 2142 ++++++++++++++++- tests/models/test_qwen3_5_registry.py | 6 + .../api_server/test_inference_endpoints.py | 58 + tests/server/backend/test_remote_backend.py | 28 + .../server/orchestrator/test_orchestrator.py | 12 +- .../test_orchestrator_client_communication.py | 4 +- .../weight_sync/test_endpoint_routing.py | 36 +- .../weight_sync/test_fp8_quantization.py | 468 ++++ .../server/weight_sync/test_handler_config.py | 56 + .../server/weight_sync/test_p2p_async_api.py | 193 ++ .../weight_sync/test_p2p_backend_protocol.py | 1377 +++++++++++ .../weight_sync/test_p2p_trainer_ib_device.py | 140 ++ 31 files changed, 6616 insertions(+), 156 deletions(-) create mode 100644 src/xorl/server/weight_sync/backends/p2p.py create mode 100644 tests/server/backend/test_remote_backend.py create mode 100644 tests/server/weight_sync/test_fp8_quantization.py create mode 100644 tests/server/weight_sync/test_handler_config.py create mode 100644 tests/server/weight_sync/test_p2p_async_api.py create mode 100644 tests/server/weight_sync/test_p2p_backend_protocol.py create mode 100644 tests/server/weight_sync/test_p2p_trainer_ib_device.py 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 c2060434..5d1cc652 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -31,6 +31,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", 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_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/server/api_server/api_types.py b/src/xorl/server/api_server/api_types.py index d1c62786..7015ad7e 100644 --- a/src/xorl/server/api_server/api_types.py +++ b/src/xorl/server/api_server/api_types.py @@ -844,6 +844,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" ) diff --git a/src/xorl/server/api_server/inference_endpoints.py b/src/xorl/server/api_server/inference_endpoints.py index a3156695..2d799c86 100644 --- a/src/xorl/server/api_server/inference_endpoints.py +++ b/src/xorl/server/api_server/inference_endpoints.py @@ -149,6 +149,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, ), ) @@ -551,6 +552,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, ) diff --git a/src/xorl/server/api_server/orchestrator_client.py b/src/xorl/server/api_server/orchestrator_client.py index 993c7599..7482d4e7 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 @@ -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/backend/remote.py b/src/xorl/server/backend/remote.py index 7f4870b2..e577a62f 100644 --- a/src/xorl/server/backend/remote.py +++ b/src/xorl/server/backend/remote.py @@ -355,7 +355,7 @@ async def sync_inference_weights( quantization=quantization, ), request_id=request_id, - timeout=600.0, + timeout=self.operation_timeout, ) async def register_adapter(self, model_id="default", lr=1e-5, request_id=None): diff --git a/src/xorl/server/orchestrator/orchestrator.py b/src/xorl/server/orchestrator/orchestrator.py index 33999b89..f7379e2d 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) β”‚ β”‚ β”‚ β”‚ ↓ ↑ β”‚ β”‚ @@ -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: diff --git a/src/xorl/server/orchestrator/request_processor.py b/src/xorl/server/orchestrator/request_processor.py index 60de2448..14fb7e66 100644 --- a/src/xorl/server/orchestrator/request_processor.py +++ b/src/xorl/server/orchestrator/request_processor.py @@ -580,6 +580,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), } diff --git a/src/xorl/server/protocol/api_orchestrator.py b/src/xorl/server/protocol/api_orchestrator.py index 310b4921..5846fee3 100644 --- a/src/xorl/server/protocol/api_orchestrator.py +++ b/src/xorl/server/protocol/api_orchestrator.py @@ -430,6 +430,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 +446,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/server_arguments.py b/src/xorl/server/server_arguments.py index e72fa973..ce393490 100644 --- a/src/xorl/server/server_arguments.py +++ b/src/xorl/server/server_arguments.py @@ -533,11 +533,13 @@ class ServerArguments: # Inference Weight Sync Configuration # ======================================================================== - sync_inference_method: Literal["nccl_broadcast"] = field( + sync_inference_method: Literal["nccl_broadcast", "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); " + "'p2p' (RDMA one-sided writes via Mooncake TransferEngine into SGLang's " + "registered param memory; requires --enable-rdma-weight-updates on the SGLang side)" }, ) diff --git a/src/xorl/server/utils/zmq_channels.py b/src/xorl/server/utils/zmq_channels.py index 349d92f6..b724ecd0 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 (connect, send) - SyncDealerChannel: Sync DEALER socket (connect, poll, recv) -- AsyncPullChannel: Async PULL socket (connect, poll, recv) +- AsyncPullChannel: Async PULL socket (bind, 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..2d7aee47 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,105 @@ 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. + +P2P tuning options: + +- `P2P_SYNC_QUANTIZATION='{"quant_method":"fp8","fmt":"e4m3","weight_block_size":[128,128]}'`: + quantizes projection weights on the trainer side and transfers FP8 weights + plus `weight_scale_inv` tensors to an FP8 SGLang receiver. +- `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. +- `XORL_P2P_BACKEND_CACHE=1`: cache P2P receiver locators and backend state + across sync calls. This is enabled by default. +- `XORL_WEIGHT_SYNC_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_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_CPU_SCRATCH_POOL_BYTES`: CPU pinned staging pool size. Keep this + above the largest staged P2P bucket; the default is 4 GiB. + +## 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 +322,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..32de8fc8 100644 --- a/src/xorl/server/weight_sync/backends/__init__.py +++ b/src/xorl/server/weight_sync/backends/__init__.py @@ -29,4 +29,8 @@ 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 == "p2p": + from .p2p import P2PTransportBackend # noqa: PLC0415 + + return P2PTransportBackend(config, **kwargs) + raise ValueError(f"Unknown weight sync backend: {method!r}. Supported: 'nccl_broadcast', '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 ebd2962f..63ad675b 100644 --- a/src/xorl/server/weight_sync/backends/nccl_broadcast.py +++ b/src/xorl/server/weight_sync/backends/nccl_broadcast.py @@ -890,7 +890,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/p2p.py b/src/xorl/server/weight_sync/backends/p2p.py new file mode 100644 index 00000000..a5f04ac4 --- /dev/null +++ b/src/xorl/server/weight_sync/backends/p2p.py @@ -0,0 +1,1945 @@ +"""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 logging +import os +import socket +import time +from concurrent.futures import Future, ThreadPoolExecutor +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] + debug_entries: List[_TransferDebugEntry] + + +@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 + + +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 _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 _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). +_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 _format_transfer_debug(debug_entries: List[_TransferDebugEntry]) -> str: + if not debug_entries: + return "transfer_debug=[]" + + parts: List[str] = [] + for entry in debug_entries[:6]: + 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 len(debug_entries) > 6: + parts.append(f"... {len(debug_entries) - 6} more") + return "transfer_debug=[" + "; ".join(parts) + "]" + + +def _chunk_sizes( + by_session: Dict[str, Tuple[List[int], List[int], List[int], List[List[_TransferDebugEntry]]]], + session_id: str, + i: int, + end: int, +) -> List[int]: + return by_session[session_id][2][i:end] + + +def _chunk_debug_entries( + by_session: Dict[str, Tuple[List[int], List[int], List[int], List[List[_TransferDebugEntry]]]], + session_id: str, + i: int, + end: int, +) -> List[_TransferDebugEntry]: + debug_lists = by_session[session_id][3][i:end] + return [entry for entries in debug_lists for entry in entries] + + +def _run_sync_transfer_items( + *, + engine_wrapper: Any, + by_session: Dict[str, Tuple[List[int], List[int], List[int], List[List[_TransferDebugEntry]]]], + 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_entries(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], List[List[_TransferDebugEntry]]]], + 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_entries(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_entries(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_entries(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, _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: + for session_id, triples in small_session_data.items(): + t_session = time.perf_counter() + for src_ptr, peer_ptr, nbytes, debug_entry in triples: + total_bytes += nbytes + session_bytes[session_id] = session_bytes.get(session_id, 0) + nbytes + num_buffers += 1 + ret = engine_wrapper.batch_transfer_sync(session_id, [src_ptr], [peer_ptr], [nbytes]) + if ret < 0: + raise RuntimeError( + f"[P2P] small-entries transfer to {session_id} " + f"failed: ret={ret} (nbytes={nbytes}, bucket {bucket_idx}, " + f"{_format_transfer_debug([debug_entry])}, " + f"session_info={session_debug_info.get(session_id)})" + ) + 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], List[List[_TransferDebugEntry]]]], + small_session_data: Dict[str, List[Tuple[int, int, int, _TransferDebugEntry]]], + session_debug_info: Dict[str, Dict[str, Any]], + small_register_ptrs: List[int], + small_register_lens: List[int], + chunk: int, + timing: _BucketTiming, + bucket_idx: int, + slice_holds: List[torch.Tensor], + src_view_holds: List[torch.Tensor], +) -> 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 os.environ.get("XORL_P2P_USE_ASYNC_API", "0") == "1": + _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 + 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). Bumped via XORL_P2P_NUM_POOLS; + # combined with XORL_P2P_MOONCAKE_WORKERS=N this gives N-way + # concurrent Mooncake calls per rank, hiding per-call latency + # on medium-speed nodes. + 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._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] = [] + # Stable group name passed back in /complete_weights_update. + self._group_name = config.group_name + self._hostname: Optional[str] = be_cfg.get("hostname") + 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)) + # 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)) + + # ------------------------------------------------------------------ + # 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 adopt_prepared_state( + self, + tensor_map: Dict[str, List[Dict[str, Any]]], + receiver_session_ids: List[str], + ) -> 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) + self._session_debug_info = { + sid: {"session_id": sid, "adopted_from_rank0": True} for sid in self._receiver_session_ids + } + 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 + has_cached_state = bool(self._tensor_map and self._receiver_session_ids) + cached_states: List[bool] = [False] * self._world_size + try: + dist.all_gather_object(cached_states, has_cached_state) + except Exception as e: + logger.warning(f"[P2P] cached-state all_gather failed; using full prepare: {e}") + cached_states = [False] * self._world_size + self._prefer_cached_prepare = all(cached_states) + + payload: List[Any] = [None] + if is_rank0: + ok = self._initialize_single_sender() + if ok: + if self._last_prepare_returned_tensor_map: + 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 the broadcast. payload[0] is None on + # failure so non-zero ranks can short-circuit cleanly. + dist.broadcast_object_list(payload, src=0) + 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 + else: + logger.error(f"[P2P] unknown initialize payload kind: {kind!r}") + local_ok = False + + init_results: List[bool] = [False] * self._world_size + try: + dist.all_gather_object(init_results, bool(local_ok)) + 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._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. + 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 + num_endpoints = len(cfg.endpoints) + while True: + if request_cached_prepare: + merged_tensor_map: Dict[str, List[Dict[str, Any]]] = { + name: [dict(loc) for loc in locators] for name, locators in self._tensor_map.items() + } + else: + merged_tensor_map = {} + merged_receiver_infos: List[Dict[str, Any]] = [] + if request_cached_prepare: + for sid in self._receiver_session_ids: + cached_info = dict(self._session_debug_info.get(str(sid), {"session_id": sid})) + cached_info.setdefault("session_id", sid) + merged_receiver_infos.append(cached_info) + + restart_full_prepare = False + for ep_idx, ep in enumerate(cfg.endpoints): + 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 request_cached_prepare: + payload["p2p_return_tensor_map"] = False + try: + resp = requests.post(url, json=payload, timeout=_prepare_timeout_seconds()) + except requests.RequestException as e: + logger.error(f"[P2P] /prepare_weights_update to {ep.host}:{ep.port} failed: {e}") + return False + if request_cached_prepare and resp.status_code in (400, 422): + logger.warning( + f"[P2P] cached prepare was rejected by {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 resp.status_code != 200: + logger.error( + f"[P2P] /prepare_weights_update returned {resp.status_code} " + f"from {ep.host}:{ep.port}: {resp.text}" + ) + return False + body = resp.json() + if not body.get("success", False): + logger.error(f"[P2P] prepare failed at {ep.host}:{ep.port}: {body.get('message')}") + return False + + 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 + 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 + + 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})" + ) + 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") + 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 + 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: + 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 = [e for e in entries if e.src_view.is_cuda and e.nbytes < self._cpu_pool_min_bytes] + large = [e for e in entries if (not e.src_view.is_cuda) or e.nbytes >= self._cpu_pool_min_bytes] + 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], List[List[_TransferDebugEntry]]]] = {} + scratch_offset_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") + 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.append( + _StagedTransfer( + src_ptr=src_ptr, + peer_ptr=e.peer_ptr, + nbytes=e.nbytes, + memory_handle=_locator_memory_handle(e.loc), + debug_entries=[_transfer_debug_entry(e.name, e.loc, e.nbytes)], + ) + ) + scratch_offset_bytes += e.nbytes + + # 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: List[List[_TransferDebugEntry]] = [] + 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 + debug_entries[-1].extend(staged_entry.debug_entries) + 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) + debug_entries.append(list(staged_entry.debug_entries)) + total_post_coalesce += len(src_ptrs) + by_session[session_id] = (src_ptrs, peer_ptrs, lens, debug_entries) + + if 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)" + ) + + # 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, _TransferDebugEntry]]] = {} + small_register_ptrs: List[int] = [] + small_register_lens: List[int] = [] + if pending_small: + small_intervals: List[Tuple[int, int]] = [] + for session_id, entries in pending_small.items(): + triples: List[Tuple[int, int, int, _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()) + small_intervals.append((s_start, s_end)) + triples.append( + (int(sv.data_ptr()), e.peer_ptr, e.nbytes, _transfer_debug_entry(e.name, e.loc, e.nbytes)) + ) + small_session_data[session_id] = triples + small_segs = self._intervals_per_cuda_segment(small_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() + + chunk = max(1, int(os.environ.get("XORL_P2P_MOONCAKE_TRANSFER_CHUNK", "1"))) + + 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=chunk, + timing=timing, + bucket_idx=bucket_idx, + slice_holds=slice_holds, + src_view_holds=src_view_holds, + ) + 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: + return frozenset(range(self._world_size)) + return frozenset({0}) + + @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 = 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 _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}") + + @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() + 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. + """ + # Prefer an FQDN that the inference side can route to. Fall back to the + # first non-loopback IPv4 address. + try: + host = socket.getfqdn() + if host and host != "localhost": + return host + except Exception: + pass + return socket.gethostbyname(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..572dbbee 100644 --- a/src/xorl/server/weight_sync/handler.py +++ b/src/xorl/server/weight_sync/handler.py @@ -26,8 +26,11 @@ This keeps GPU memory to ~1-2 layers (one unsharding + one broadcasting). """ +import atexit import logging +import os import time +from concurrent.futures import Future, ThreadPoolExecutor from typing import Any, Dict, List, Optional, Tuple import torch @@ -54,8 +57,29 @@ 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; the env var remains an explicit override.""" + default = _DEFAULT_P2P_MOE_BUCKET_BYTES if sync_method == "p2p" else _DEFAULT_MOE_BUCKET_BYTES + return _env_int("XORL_WEIGHT_SYNC_BUCKET_BYTES", default) def _prod(shape) -> int: @@ -66,6 +90,141 @@ 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 + + +def _atexit_destroy_cached_backend() -> None: + global _cached_p2p_backend, _cached_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 + + +atexit.register(_atexit_destroy_cached_backend) + + class WeightSyncHandler: """Handles weight synchronization between training and inference endpoints.""" @@ -73,6 +232,176 @@ 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]] = {} + + 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 # ======================================================================== # Main entry point @@ -83,7 +412,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 +485,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 +522,37 @@ 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 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 + transport_cfg = TransportConfig( endpoints=[ EndpointConfig( @@ -204,9 +569,58 @@ def _sync_weights( device=device, training_world_size=self.world_size, training_rank=self.rank, + backend_config=_backend_config, ) - backend = create_backend(sync_method, transport_cfg) + # 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 + cache_enabled = os.environ.get("XORL_P2P_BACKEND_CACHE", "1") == "1" and sync_method == "p2p" + 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, v) + for k, v in (_backend_config or {}).items() + if k not in {"flush_cache", "weight_version"} and isinstance(v, (str, int, bool, float)) + ) + ), + ) + if ( + cache_enabled + 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 + 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 +630,35 @@ 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: 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") # 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 +671,21 @@ 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. + # Default off β€” flip via XORL_WEIGHT_SYNC_BATCH_MOE=1. + 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._reset_fp8_cpu_workspace_usage() # Build ordered list of FSDP modules to process modules_to_sync: List[Tuple[str, FSDPModule]] = [] @@ -295,31 +733,83 @@ 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, + 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() @@ -334,6 +824,7 @@ def _sync_weights( ep_moe_prefixes.add(p) if _stage_leader: + 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}" @@ -345,14 +836,21 @@ def _sync_weights( skip_moe_prefixes=ep_moe_prefixes, ) 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: + t_phase = time.perf_counter() current_buffer = self._unfuse_for_inference( current_buffer, model, @@ -361,46 +859,109 @@ def _sync_weights( 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", ) + _add_rank_phase("unfuse_quantize_s", t_phase) + if _ws_timings: + _t_unfuse = time.perf_counter() logger.info(f"Rank 0: [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 - - # Stage 0 MoE handling (unchanged) + 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 and self.rank in backend.sender_ranks + ) 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 +989,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 +1008,148 @@ 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: + # 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 +1158,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,6 +1336,8 @@ def _collect_ep_moe_data( fsdp_mod, mod_name: str, ps, + skip_clone: bool = False, + phase_s: Optional[Dict[str, float]] = None, ) -> List[Dict[str, Any]]: """Collect local EP-sharded MoE expert data during unshard phase. @@ -678,16 +1396,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 +1436,488 @@ 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, + ) + + 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"] + + 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 + + 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) + # 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() + 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() + 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() + 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, @@ -835,6 +2049,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 +2070,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, @@ -955,6 +2171,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 +2190,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, @@ -1104,10 +2322,704 @@ 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 _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 + ) + + return "_proj.weight" in name or name.endswith("fused_qkv_a_proj_with_mqa.weight") + + @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 +3033,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 @@ -1335,7 +3235,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)) @@ -1371,6 +3272,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 +3318,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, diff --git a/tests/models/test_qwen3_5_registry.py b/tests/models/test_qwen3_5_registry.py index 5fe1bf4e..d2c3d9e7 100644 --- a/tests/models/test_qwen3_5_registry.py +++ b/tests/models/test_qwen3_5_registry.py @@ -87,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) @@ -96,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/server/api_server/test_inference_endpoints.py b/tests/server/api_server/test_inference_endpoints.py index 602fff68..4d7d0707 100644 --- a/tests/server/api_server/test_inference_endpoints.py +++ b/tests/server/api_server/test_inference_endpoints.py @@ -107,6 +107,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 +132,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/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/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..6f469a59 --- /dev/null +++ b/tests/server/weight_sync/test_fp8_quantization.py @@ -0,0 +1,468 @@ +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 + + +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..2add0946 --- /dev/null +++ b/tests/server/weight_sync/test_handler_config.py @@ -0,0 +1,56 @@ +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, +) + + +def test_moe_bucket_default_is_backend_specific(monkeypatch): + 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_moe_bucket_env_override_is_explicit(monkeypatch): + 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_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_p2p_async_api.py b/tests/server/weight_sync/test_p2p_async_api.py new file mode 100644 index 00000000..69a7f499 --- /dev/null +++ b/tests/server/weight_sync/test_p2p_async_api.py @@ -0,0 +1,193 @@ +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, + timing=timing, + bucket_idx=1, + slice_holds=[], + src_view_holds=[], + ) + + 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, + timing=timing, + bucket_idx=1, + slice_holds=[], + src_view_holds=[], + ) + + 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, + timing=timing, + bucket_idx=1, + slice_holds=[], + src_view_holds=[], + ) + + 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, + timing=_BucketTiming(), + bucket_idx=7, + slice_holds=[], + src_view_holds=[], + ) + + +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..db19ab25 --- /dev/null +++ b/tests/server/weight_sync/test_p2p_backend_protocol.py @@ -0,0 +1,1377 @@ +"""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 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 + + +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 _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_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): + 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._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): + 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_restarts_all_endpoints_when_later_endpoint_rejects_flag(self): + 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 + ep0_full_locator = { + **_hf_locator( + tp_rank=0, + full_shape=[128, 64], + slc=[[0, 64], [0, 64]], + ptr=0x1000, + nbytes=64 * 64 * 2, + session_id="recv-0:7000", + ), + "hf_name": param_name, + } + 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"}], + }, + ) + ep0_full_response = _FakeResponse( + 200, + { + "success": True, + "message": "full", + "tensor_map": {param_name: [ep0_full_locator]}, + "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, {}), ep0_full_response, ep1_full_response], + ) as posted: + ok = backend.initialize() + + assert ok is True + assert posted.call_count == 4 + 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-0:5000/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"] + assert "p2p_return_tensor_map" not in posted.call_args_list[3].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": 0x1000, + "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"}} + + +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_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_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") + 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 + + +class TestP2PDirectEPTransfer: + def _make_multi_sender_backend( + self, *, rank_index: int, world_size: int, rank_filter + ) -> Tuple[P2PTransportBackend, _FakeMooncakeEngine]: + 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={ + "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, + }, + ) + 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}) + + 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): + 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_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_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}, + ] From f995932c629aab6456a8df8c20f43c5a5f29f14f Mon Sep 17 00:00:00 2001 From: Ashwinee Panda Date: Thu, 14 May 2026 16:52:41 -0700 Subject: [PATCH 33/49] feat(checkpoint): grouped HF weight load + DCP convert script * feat(checkpoint): grouped HF weight load + DCP convert script Adds two complementary loaders for very large MoE checkpoints (1T+ params) that the existing 'broadcast' and 'all_ranks' modes don't handle well: - 'grouped' load mode (xorl.models.module_utils.grouped_load_weights) uses two fan-out tiers: one reader per node broadcasts dense/shared tensors inside a node-local process group, and one reader per EP-FSDP group leader scatters expert tensors to that group's peers. Each shard is read from disk once per node (not once per rank), and expert reads are filtered to the [ep_start, ep_end) slice via safetensors load_filtered + the model's get_skip_key_fn. - 'skip' load mode is a sibling that bypasses HF loading entirely, intended for resuming from a DCP checkpoint (load_checkpoint_path). __post_init__ rejects skip without a checkpoint path so the misuse fails fast. - scripts/convert_checkpoint.py is a standalone CLI that builds the model under 'grouped' load, then writes a per-rank-sharded DCP via DistributedCheckpointer.save. The output directory becomes a load_checkpoint_path that subsequent 'skip'-mode runs can consume. Supporting pieces: - ModelState.reference_state_dict() for direct safetensors export from live (sharded) tensors without DCP-style materialization. - save_model_weights_distributed() for the safetensors export path. - ExpertWeightBuffer GPU fast-path now skips the CPU bounce when the source tensor is already on the target device (e.g. inline dequant). - DistributedCheckpointer.load now no-ops a missing extra_state file instead of raising, so load_checkpoint_path can point at a convert output that has no trainer extra_state yet. EP-aware load_filtered already existed on main for 'all_ranks' and 'broadcast'; this PR adds 'grouped' as a third caller and the per-node + per-EP-FSDP fan-out groups under it. Real-world numbers on ~600 GB Kimi-K2.5 DCP (warm-cache load): 16 pods, ep=16/efsdp=8: ~57s 24 pods, ep=24/efsdp=8: ~85s * Fix grouped loader expert key routing * Fix grouped loader expert key routing * Add grouped load materialization helper * Fix materialize grouped load mesh defaults * feat(checkpoint): default to grouped weight loading * refactor(checkpoint): remove broadcast load mode * chore: apply pre-commit formatting --- .../content/docs/config-reference/local.md | 2 +- .../content/docs/config-reference/server.md | 2 +- .../docs/parallelism/data_parallelism.mdx | 21 +- .../docs/parallelism/expert_parallelism.mdx | 4 +- .../docs/parallelism/pipeline_parallelism.mdx | 2 +- .../docs/parallelism/tensor_parallelism.mdx | 12 +- .../content/docs/training/local_training.mdx | 6 +- docs/src/content/docs/training/system.mdx | 4 +- .../coderforge/configs/qwen3_5_35b_a3b.yaml | 2 +- .../configs/qwen3_5_35b_a3b_muon.yaml | 2 +- .../local/dummy/configs/full/llama3_8b.yaml | 2 +- .../dummy/configs/full/llama3_8b_pp2.yaml | 2 +- .../dummy/configs/full/llama3_8b_tp4.yaml | 2 +- .../full/qwen3_30b_a3b_pp2_ep4_cp4_muon.yaml | 2 +- .../full/qwen3_30b_a3b_pp2_ep4_muon.yaml | 2 +- .../local/dummy/configs/full/qwen3_32b.yaml | 2 +- .../dummy/configs/full/qwen3_5_35b_a3b.yaml | 2 +- .../local/dummy/configs/full/qwen3_8b.yaml | 2 +- .../dummy/configs/full/qwen3_8b_cp1_sp8.yaml | 2 +- .../dummy/configs/full/qwen3_8b_cp2_sp4.yaml | 2 +- .../dummy/configs/full/qwen3_8b_cp8_sp1.yaml | 2 +- .../dummy/configs/full/qwen3_8b_muon.yaml | 2 +- .../dummy/configs/full/qwen3_8b_pp2.yaml | 2 +- .../configs/full/qwen3_8b_sp4_full_128k.yaml | 2 +- .../configs/full/qwen3_8b_tp4_compile.yaml | 2 +- .../dummy/configs/lora/llama3_8b_lora.yaml | 2 +- .../dummy/configs/lora/qwen3_8b_lora.yaml | 2 +- .../configs/lora/qwen3_8b_lora_cp2_sp2.yaml | 2 +- .../dummy/configs/lora/qwen3_8b_lora_cp4.yaml | 2 +- .../configs/qlora/llama3_8b_qlora_nvfp4.yaml | 2 +- .../qlora/qwen3_32b_fp8_prequant_qlora.yaml | 2 +- .../qlora/qwen3_32b_nvfp4_prequant_qlora.yaml | 2 +- .../configs/qlora/qwen3_32b_qlora_nvfp4.yaml | 2 +- .../qlora/qwen3_8b_qlora_block_fp8.yaml | 2 +- .../configs/qlora/qwen3_8b_qlora_nvfp4.yaml | 2 +- .../qlora/qwen3_8b_qlora_nvfp4_pp2.yaml | 2 +- .../qlora/qwen3_8b_qlora_nvfp4_requant.yaml | 2 +- .../qlora/qwen3_8b_qlora_nvfp4_sp4.yaml | 2 +- src/xorl/arguments.py | 18 +- src/xorl/checkpoint/checkpointer.py | 11 + src/xorl/distributed/torch_parallelize.py | 31 +- src/xorl/models/module_utils.py | 54 ++- src/xorl/qlora/utils.py | 35 +- src/xorl/server/runner/model_runner.py | 2 +- src/xorl/server/server_arguments.py | 8 +- src/xorl/trainers/model_builder.py | 4 +- tests/models/test_model_state.py | 26 ++ tests/models/test_module_utils_broadcast.py | 310 +++++++++++++++++- tests/server/test_server_arguments.py | 23 ++ tests/test_arguments.py | 43 +++ 50 files changed, 580 insertions(+), 98 deletions(-) diff --git a/docs/src/content/docs/config-reference/local.md b/docs/src/content/docs/config-reference/local.md index 5bead0b1..12ba3a83 100644 --- a/docs/src/content/docs/config-reference/local.md +++ b/docs/src/content/docs/config-reference/local.md @@ -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..c840d6a7 100644 --- a/docs/src/content/docs/config-reference/server.md +++ b/docs/src/content/docs/config-reference/server.md @@ -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). | --- 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..80cc4903 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 @@ -837,7 +837,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/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/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/src/xorl/arguments.py b/src/xorl/arguments.py index 14cbf65a..e9213923 100644 --- a/src/xorl/arguments.py +++ b/src/xorl/arguments.py @@ -821,10 +821,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", "grouped", "skip"] = field( - default="broadcast", + load_weights_mode: Literal["all_ranks", "grouped", "skip"] = field( + default="grouped", metadata={ - "help": "Weight loading mode. 'broadcast': global rank0 reads weights and broadcasts to other ranks. 'all_ranks': every rank reads weights from disk independently. 'grouped': one reader per EP-FSDP group reads and broadcasts within that group. 'skip': skip HF weight loading (use with load_checkpoint_path for DCP)." + "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( @@ -1109,6 +1116,11 @@ 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 " diff --git a/src/xorl/checkpoint/checkpointer.py b/src/xorl/checkpoint/checkpointer.py index 83f69bb0..7c5670ad 100644 --- a/src/xorl/checkpoint/checkpointer.py +++ b/src/xorl/checkpoint/checkpointer.py @@ -636,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, diff --git a/src/xorl/distributed/torch_parallelize.py b/src/xorl/distributed/torch_parallelize.py index cc92af42..92b6af70 100644 --- a/src/xorl/distributed/torch_parallelize.py +++ b/src/xorl/distributed/torch_parallelize.py @@ -19,7 +19,7 @@ pipeline_module_split, ) from xorl.lora import LoraLinear -from xorl.models import all_ranks_load_weights, grouped_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 @@ -46,18 +46,27 @@ def _load_model_weights( 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 - elif 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=dtensor_factory) - elif load_weights_mode == "grouped": - logger.info_rank0("Loading model weights with one reader per EP-FSDP group...") + 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) - else: + 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 = { @@ -424,7 +433,7 @@ def _experts_shard_placement_fn(param): 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: @@ -438,7 +447,7 @@ 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") + load_weights_mode = kwargs.get("load_weights_mode", "grouped") _load_model_weights( model, weights_path, @@ -640,7 +649,7 @@ 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") + load_weights_mode = kwargs.get("load_weights_mode", "grouped") _load_model_weights( model_part, weights_path, @@ -784,7 +793,7 @@ 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") + load_weights_mode = kwargs.get("load_weights_mode", "grouped") _load_model_weights( model, weights_path, diff --git a/src/xorl/models/module_utils.py b/src/xorl/models/module_utils.py index e97cb1d8..dd9e3096 100644 --- a/src/xorl/models/module_utils.py +++ b/src/xorl/models/module_utils.py @@ -346,12 +346,29 @@ def _normalize_checkpoint_key_for_filter(key: str) -> Optional[str]: 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 + 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: @@ -1784,6 +1801,12 @@ def grouped_load_weights( 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()} @@ -1796,14 +1819,6 @@ def grouped_load_weights( if hasattr(model, "get_parallel_plan"): parallel_plan = model.get_parallel_plan() - _ps = get_parallel_state() - fanout_group = _get_grouped_weight_load_group(_ps) - if fanout_group is None: - logger.warning_once( - "Grouped weight loading requires EP/FSDP groups; falling back to rank0_load_and_broadcast_weights." - ) - return rank0_load_and_broadcast_weights(model, weights_path, init_device, dtensor_factory) - fanout_ranks = dist.get_process_group_ranks(fanout_group) fanout_src = fanout_ranks[0] dense_group = _get_grouped_dense_weight_load_group() @@ -1851,6 +1866,7 @@ def grouped_load_weights( 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: @@ -2030,12 +2046,24 @@ def _broadcast_queue_and_dispatch( 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_checkpoint_key_for_filter(key) - return normalized is None or _is_checkpoint_expert_key(normalized) + 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_checkpoint_key_for_filter(key) + 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: @@ -2098,7 +2126,7 @@ def _should_skip_grouped_expert_key(key: str) -> bool: if is_group_leader: state_dict, skipped_keys = next(expert_prefetched) for skipped_key in skipped_keys: - normalized = _normalize_checkpoint_key_for_filter(skipped_key) + 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): 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/runner/model_runner.py b/src/xorl/server/runner/model_runner.py index b4206883..c1cd9901 100644 --- a/src/xorl/server/runner/model_runner.py +++ b/src/xorl/server/runner/model_runner.py @@ -492,7 +492,7 @@ 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, diff --git a/src/xorl/server/server_arguments.py b/src/xorl/server/server_arguments.py index ce393490..cbce8994 100644 --- a/src/xorl/server/server_arguments.py +++ b/src/xorl/server/server_arguments.py @@ -229,7 +229,8 @@ 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( @@ -566,6 +567,11 @@ def __post_init__(self): # the launcher can still use them via engine_connect_host + worker_bind_port pass + 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 " diff --git a/src/xorl/trainers/model_builder.py b/src/xorl/trainers/model_builder.py index 6459f38e..362e08e3 100644 --- a/src/xorl/trainers/model_builder.py +++ b/src/xorl/trainers/model_builder.py @@ -136,7 +136,7 @@ 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, @@ -477,7 +477,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/tests/models/test_model_state.py b/tests/models/test_model_state.py index ef7e728f..4ccd5934 100644 --- a/tests/models/test_model_state.py +++ b/tests/models/test_model_state.py @@ -31,3 +31,29 @@ def test_reference_state_dict_bypasses_dcp_state_dict_and_skips_nonpersistent_bu 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 5c03de39..584d4f2a 100644 --- a/tests/models/test_module_utils_broadcast.py +++ b/tests/models/test_module_utils_broadcast.py @@ -447,6 +447,28 @@ def fake_try_load_state_dict_local(weights_path, **kwargs): 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 = [] @@ -638,6 +660,290 @@ def fail_prefetch(*args, **kwargs): ) +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 = [] @@ -659,6 +965,8 @@ def test_grouped_load_weights_falls_back_without_ep_group(monkeypatch): lambda *args, **kwargs: called.append((args, kwargs)), ) - module_utils.grouped_load_weights(_DummyModel(), "dummy-weights", init_device="cpu") + 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/server/test_server_arguments.py b/tests/server/test_server_arguments.py index c09914a0..7ca35b29 100644 --- a/tests/server/test_server_arguments.py +++ b/tests/server/test_server_arguments.py @@ -27,7 +27,9 @@ 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): @@ -53,6 +55,27 @@ def test_load_server_arguments_threads_distsignsgd_through_nested_config(tmp_pat assert args.to_config_dict()["train"]["optimizer"] == "distsignsgd" +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_threads_muon_gram_newton_schulz_through_nested_config(tmp_path): config_path = tmp_path / "server_config.yaml" config_path.write_text( diff --git a/tests/test_arguments.py b/tests/test_arguments.py index 157fef54..c973111e 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,7 @@ 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): @@ -125,3 +127,44 @@ def test_parse_args_wires_muon_kwargs_from_yaml(tmp_path, monkeypatch): 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 From 22b24f39d8a73555e6a118bc46375554cd6646c0 Mon Sep 17 00:00:00 2001 From: Ashwinee Panda Date: Thu, 14 May 2026 23:11:00 -0700 Subject: [PATCH 34/49] feat: multi-LoRA * Start multi-adapter LoRA refactor * Add configurable adapter state restore mode * Add multi-adapter LoRA parity harness * Refactor multi-adapter LoRA server optimizer wiring * Support heterogeneous multi-adapter LoRA sessions * Sync local multi-adapter test and import fixes * Fix server session lifecycle semantics * Add K8s heterogeneous multi-adapter workflow * Add k8s multi-adapter LoRA stress and SGL sync validation * Fix multi-adapter checkpoint and MoE sync edge cases * Fix session checkpoint lifecycle and legacy optim step mapping * Fix adapter eviction and coordinated checkpoint restore * Fix routing replay transport and sampling session tracking * Fix LoRA session save and registration failure paths * Harden adapter restore and checkpoint save failures * Fix weights-only restore LR and kill-session checkpoint paths * Tighten multi-adapter LoRA checkpoint restore validation * Harden multi-adapter LoRA restore coordination * Fix non-packed password memorization evaluation * Harden sampler adapter checkpoint compatibility * Fix merged API type test expectations * Fix pre-commit after main merge * Roll back failed adapter checkpoint loads * Split multi-adapter LoRA foundation * fix: drop stale worker_port reintroduction and update OptimStepRequest test Main removed worker_port from InferenceEndpoint in. This branch predated that and was re-adding it as a required field, breaking add_inference_endpoint and SyncInferenceWeights tests. Drop the worker_port re-addition. Also update test_request_creation_and_serialization to match the new OptimStepRequest API where adam_params defaults to None and learning_rate is the top-level field. * fix: preserve Tinker API compatibility in multi-LoRA foundation * fix: address multi-LoRA foundation review feedback * fix: address PR 265 review regressions * Split multi-adapter LoRA runtime rank substrate * fix: address runtime rank LoRA review feedback * Split multi-adapter LoRA adapter manager core * fix: address adapter manager review feedback * Split multi-adapter LoRA worker coordination * Fix PR 268 worker session coordination regressions * Split multi-adapter LoRA API session lifecycle * Fix API session lifecycle edge cases * Split multi-adapter LoRA checkpoint restore hardening * fix: harden multi-adapter LoRA checkpoint export * Split multi-adapter LoRA sampler inference tracking * fix sampler inference endpoint tracking * fix: route sampler adapter calls to worker port * Split multi-adapter LoRA validation docs and experiments * Fix multi-LoRA stack test expectations * Route register_session through adapter coordinator * Constrain multi-LoRA compare job to nccl nodes * Merge remote-tracking branch 'origin/main' into split/mal-01-foundation * fix: export runtime-rank LoRA state * fix: align adapter coordinator checkpoint loads * Fix API session refresh edge cases * fix: slice multi-adapter lora exports to session rank * fix: detect resolved MoE LoRA targets on save * Fix updated multi-LoRA stack integration * Preserve LoRA dtype for training checkpoints * Document main-merged multi-LoRA validation * Enable deterministic FA3 replay for multi-LoRA tests * Update multi-LoRA validation report * Sanitize multi-LoRA experiment manifests * Serialize IS metrics as scalars * Use namespace-qualified password sampler endpoint * Define password sampler namespace in trainer job * fix: address multi-LoRA review issues * fix: support Kimi tokenizer in multi-LoRA harness * Fix Kimi EP32 multi-LoRA checkpoint restore * Add server full determinism option * fix multi-adapter server review regressions * fix full-weight API compatibility regressions * fix ep-sharded adapter rank0 restore * Validate session spec for EP broadcast restores * chore: apply pre-commit formatting * fix: update register session dispatcher test --------- Co-authored-by: Qingyang Wu --- .gitignore | 3 + .../content/docs/config-reference/server.md | 1 + .../content/docs/server-training/overview.md | 2 + .../training-server/api-reference.md | 4 +- examples/server/no_robot_sft/README.md | 4 +- examples/server/no_robot_sft/run_sft.py | 6 +- .../run_password_test.py | 6 +- .../run_train_and_infer.py | 4 +- src/xorl/arguments.py | 4 + src/xorl/distributed/offloading.py | 5 +- src/xorl/lora/modules/linear.py | 34 +- src/xorl/lora/utils.py | 277 ++++++- src/xorl/models/auto.py | 2 + .../attention/backend/flash_attention.py | 10 + src/xorl/models/layers/moe/lora.py | 100 +-- src/xorl/ops/group_gemm/kernel/__init__.py | 31 +- src/xorl/ops/group_gemm/kernel/lora_utils.py | 148 ++++ src/xorl/qlora/modules/linear.py | 22 +- src/xorl/qlora/modules/moe_experts.py | 83 +- src/xorl/server/api_server/__init__.py | 6 +- src/xorl/server/api_server/api_types.py | 156 +++- src/xorl/server/api_server/endpoints.py | 335 +++++--- .../server/api_server/inference_endpoints.py | 225 ++++-- .../server/api_server/orchestrator_client.py | 2 +- src/xorl/server/api_server/server.py | 131 ++-- src/xorl/server/api_server/training_ops.py | 52 +- src/xorl/server/api_server/weights.py | 31 +- src/xorl/server/backend/base.py | 13 + src/xorl/server/backend/dummy.py | 28 +- src/xorl/server/backend/remote.py | 36 +- src/xorl/server/launcher.py | 55 +- src/xorl/server/orchestrator/orchestrator.py | 5 +- src/xorl/server/orchestrator/packing.py | 53 +- .../server/orchestrator/request_processor.py | 25 +- src/xorl/server/protocol/__init__.py | 1 + src/xorl/server/protocol/api_orchestrator.py | 1 + src/xorl/server/protocol/operations.py | 11 + .../server/protocol/orchestrator_runner.py | 1 + .../runner/adapters/adapter_coordinator.py | 723 ++++++++++++++++-- src/xorl/server/runner/adapters/manager.py | 631 ++++++++++++--- src/xorl/server/runner/checkpoint/manager.py | 193 ++++- src/xorl/server/runner/model_runner.py | 638 +++++++++++----- src/xorl/server/runner/runner_dispatcher.py | 70 +- src/xorl/server/server_arguments.py | 100 +++ src/xorl/server/session_spec.py | 442 +++++++++++ src/xorl/server/utils/zmq_channels.py | 4 +- src/xorl/server/weight_sync/handler.py | 51 +- src/xorl/trainers/model_builder.py | 2 + src/xorl/trainers/trainer.py | 1 + tests/distributed/test_offloading.py | 19 + tests/e2e/qwen3_8b/test_tflops_threshold.py | 5 +- tests/e2e/server_utils.py | 6 +- tests/models/test_moe_experts_lora.py | 37 +- tests/ops/test_attention.py | 16 + tests/ops/test_lora_utils.py | 41 + tests/server/api_server/test_api_server.py | 483 +++++++++++- tests/server/api_server/test_api_types.py | 124 ++- .../api_server/test_checkpoint_paths.py | 156 +++- .../api_server/test_inference_endpoints.py | 152 ++++ .../api_server/test_session_endpoints.py | 634 +++++++++++++++ tests/server/api_server/test_training_ops.py | 114 +++ tests/server/orchestrator/test_packing.py | 81 +- .../orchestrator/test_request_processor.py | 240 ++++++ .../server/runner/test_adapter_coordinator.py | 645 ++++++++++++++++ tests/server/runner/test_adapter_manager.py | 613 +++++++++++++++ .../test_checkpoint_manager_save_failures.py | 336 ++++++++ .../runner/test_lora_checkpoint_roundtrip.py | 94 ++- .../runner/test_model_runner_is_metrics.py | 84 ++ .../runner/test_model_runner_kill_session.py | 76 ++ .../test_model_runner_session_registry.py | 186 +++++ .../runner/test_runner_dispatcher_forward.py | 79 ++ .../test_runner_dispatcher_load_state.py | 110 +++ .../test_runner_dispatcher_session_ops.py | 164 ++++ tests/server/test_server_arguments.py | 197 ++++- .../weight_sync/test_moe_runtime_scaling.py | 78 ++ 75 files changed, 8665 insertions(+), 873 deletions(-) create mode 100644 src/xorl/ops/group_gemm/kernel/lora_utils.py create mode 100644 src/xorl/server/session_spec.py create mode 100644 tests/distributed/test_offloading.py create mode 100644 tests/ops/test_lora_utils.py create mode 100644 tests/server/api_server/test_session_endpoints.py create mode 100644 tests/server/api_server/test_training_ops.py create mode 100644 tests/server/runner/test_adapter_coordinator.py create mode 100644 tests/server/runner/test_adapter_manager.py create mode 100644 tests/server/runner/test_checkpoint_manager_save_failures.py create mode 100644 tests/server/runner/test_model_runner_is_metrics.py create mode 100644 tests/server/runner/test_model_runner_kill_session.py create mode 100644 tests/server/runner/test_model_runner_session_registry.py create mode 100644 tests/server/runner/test_runner_dispatcher_forward.py create mode 100644 tests/server/runner/test_runner_dispatcher_load_state.py create mode 100644 tests/server/runner/test_runner_dispatcher_session_ops.py create mode 100644 tests/server/weight_sync/test_moe_runtime_scaling.py diff --git a/.gitignore b/.gitignore index aaa5bd5e..673ea19e 100644 --- a/.gitignore +++ b/.gitignore @@ -83,3 +83,6 @@ 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/docs/src/content/docs/config-reference/server.md b/docs/src/content/docs/config-reference/server.md index c840d6a7..e376c166 100644 --- a/docs/src/content/docs/config-reference/server.md +++ b/docs/src/content/docs/config-reference/server.md @@ -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/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/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/src/xorl/arguments.py b/src/xorl/arguments.py index e9213923..18281bc4 100644 --- a/src/xorl/arguments.py +++ b/src/xorl/arguments.py @@ -441,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={ 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/lora/modules/linear.py b/src/xorl/lora/modules/linear.py index b7510468..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: """ @@ -195,16 +211,9 @@ def merge_weights(self) -> None: Used both for periodic merge during training (merge_lora_interval) and one-shot merge for inference. - - Precision note: we upcast ``weight`` to float32, add the fp32 delta, - then cast the sum back β€” a *single* quantization at the end. This is - strictly more faithful than the naive ``weight += delta.to(W.dtype)`` - (which rounds Ξ” per element before adding) and matters in precision- - sensitive settings like MoE top-k routing on bf16 weights. Same memory - as the naive variant β€” both land in ``self.weight.dtype``. """ with torch.no_grad(): - delta_weight = self.get_delta_weight() # fp32 + delta_weight = self.get_delta_weight() merged = self.weight.to(torch.float32) + delta_weight self.weight.data.copy_(merged.to(self.weight.dtype)) self.reset_lora_parameters() @@ -213,5 +222,6 @@ 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 a58e1b0b..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 @@ -76,6 +77,15 @@ 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) @@ -328,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 = "", @@ -349,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(): @@ -368,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 @@ -433,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 @@ -446,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. @@ -471,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. @@ -492,7 +696,9 @@ def convert_peft_lora_state_dict( # 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) + 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) @@ -506,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: @@ -524,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 @@ -540,6 +750,7 @@ def save_lora_checkpoint( lora_state_dict: Optional[Dict[str, torch.Tensor]] = None, 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. @@ -568,6 +779,9 @@ def save_lora_checkpoint( (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 @@ -586,6 +800,14 @@ def save_lora_checkpoint( # 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.""" @@ -643,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): @@ -653,7 +875,7 @@ 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 @@ -669,7 +891,7 @@ def _sglang_shared_outer_moe_weight(name: str, stacked_tensor: torch.Tensor) -> 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 = stacked_tensor.transpose(-2, -1).contiguous().to(torch.bfloat16) + 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} @@ -708,18 +930,32 @@ def _sglang_shared_outer_moe_weight(name: str, stacked_tensor: torch.Tensor) -> # 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 @@ -797,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) @@ -997,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) @@ -1005,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/auto.py b/src/xorl/models/auto.py index f2a9f919..0d4d8c77 100644 --- a/src/xorl/models/auto.py +++ b/src/xorl/models/auto.py @@ -182,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: @@ -227,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( 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/lora.py b/src/xorl/models/layers/moe/lora.py index 1550267a..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,33 +165,41 @@ 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. After merging: weight = weight + delta_weight for each active projection. Resets LoRA parameters after merge. - - Precision note: ``base`` is upcast to float32, the fp32 delta is added, - and the sum is cast back once. This is strictly more faithful than - rounding Ξ” per element before adding, and keeps unmerged-forward close - to merged-forward on MoE models where top-k routing amplifies any - per-element delta quantization error. """ with torch.no_grad(): for proj_name in ("gate_proj", "up_proj", "down_proj"): if proj_name not in self.lora_config.target_modules: continue base = getattr(self, proj_name) - delta = self._compute_proj_delta(proj_name) # fp32 + delta = self._compute_proj_delta(proj_name) merged = base.to(torch.float32) + delta base.data.copy_(merged.to(base.dtype)) self.reset_lora_parameters() @@ -226,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, @@ -234,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( @@ -271,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) @@ -283,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) @@ -346,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/ops/group_gemm/kernel/__init__.py b/src/xorl/ops/group_gemm/kernel/__init__.py index e7fbf583..c4fe4487 100644 --- a/src/xorl/ops/group_gemm/kernel/__init__.py +++ b/src/xorl/ops/group_gemm/kernel/__init__.py @@ -1,8 +1,15 @@ -import math - # Group GEMM kernels from .group_gemm import group_gemm_same_mn, group_gemm_same_nk +# LoRA utilities +from .lora_utils import ( + compute_lora_scaling, + get_lora_delta_weight_stacked, + init_lora_weights_stacked, + merge_lora_weights_stacked, + unmerge_lora_weights_stacked, +) + # MoE operations from .moe import ( expert_histogram, @@ -14,22 +21,6 @@ from .quack import quack_group_gemm_same_mn, quack_group_gemm_same_nk -def compute_lora_scaling(lora_alpha: int, r: int, use_rslora: bool = False) -> float: - """Compute the LoRA scaling factor. - - Args: - lora_alpha: LoRA alpha parameter. - r: LoRA rank. - use_rslora: Whether to use rank-stabilized LoRA scaling. - - Returns: - Scaling factor. - """ - if use_rslora: - return lora_alpha / math.sqrt(r) - return lora_alpha / r - - __all__ = [ # Group GEMM "group_gemm_same_mn", @@ -43,5 +34,9 @@ def compute_lora_scaling(lora_alpha: int, r: int, use_rslora: bool = False) -> f "moe_index_compute", "moe_scatter", # LoRA utilities + "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/ops/group_gemm/kernel/lora_utils.py b/src/xorl/ops/group_gemm/kernel/lora_utils.py new file mode 100644 index 00000000..9482ab2e --- /dev/null +++ b/src/xorl/ops/group_gemm/kernel/lora_utils.py @@ -0,0 +1,148 @@ +"""LoRA utilities for MoE implementation. + +This module provides utilities for initializing and managing LoRA weights +in the stacked tensor format used by group GEMM kernels. +""" + +import math +from typing import Optional, Tuple + +import torch +import torch.nn as nn + + +def init_lora_weights_stacked( + num_experts: int, + r: int, + in_features: int, + out_features: int, + init_method: str = "kaiming", + dtype: torch.dtype = torch.float32, + device: Optional[torch.device] = None, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Initialize stacked LoRA weights for all experts. + + Creates lora_A and lora_B tensors with appropriate initialization: + - lora_A: Kaiming uniform or Gaussian initialization + - lora_B: Zero initialization (ensures delta_W = 0 at start) + + Args: + num_experts: Number of experts + r: LoRA rank + in_features: Input feature dimension + out_features: Output feature dimension + init_method: Initialization method ("kaiming" or "gaussian") + dtype: Data type for the tensors + device: Device for the tensors + + Returns: + Tuple of (lora_A, lora_B) tensors: + - 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, 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, r, out_features] + lora_B = torch.zeros(num_experts, r, out_features, dtype=dtype, device=device) + + # Initialize lora_A + if init_method == "kaiming": + for i in range(num_experts): + # Initialize each expert's lora_A with kaiming uniform + nn.init.kaiming_uniform_(lora_A[i], a=math.sqrt(5)) + elif init_method == "gaussian": + nn.init.normal_(lora_A, std=1.0 / r) + else: + raise ValueError(f"Unknown init_method: {init_method}") + + # lora_B is already zeros + + return lora_A, lora_B + + +def compute_lora_scaling(lora_alpha: int, r: int, use_rslora: bool = False) -> float: + """Compute the LoRA scaling factor. + + Args: + lora_alpha: LoRA alpha parameter + r: LoRA rank + use_rslora: Whether to use rank-stabilized LoRA scaling + + Returns: + Scaling factor + """ + if use_rslora: + return lora_alpha / math.sqrt(r) + else: + return lora_alpha / r + + +def merge_lora_weights_stacked( + base_weight: torch.Tensor, + lora_A: torch.Tensor, + lora_B: torch.Tensor, + scaling: float, +) -> torch.Tensor: + """Merge LoRA weights into base weights. + + Computes: W' = W + A @ B * scaling + + Args: + 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, in_features, out_features] + """ + # 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 + + +def unmerge_lora_weights_stacked( + merged_weight: torch.Tensor, + lora_A: torch.Tensor, + lora_B: torch.Tensor, + scaling: float, +) -> torch.Tensor: + """Unmerge LoRA weights from merged weights. + + Computes: W = W' - A @ B * scaling + + Args: + 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, in_features, out_features] + """ + delta_weight = torch.bmm(lora_A, lora_B) * scaling + return merged_weight - delta_weight + + +def get_lora_delta_weight_stacked( + lora_A: torch.Tensor, + lora_B: torch.Tensor, + scaling: float, +) -> torch.Tensor: + """Compute the LoRA weight delta. + + Computes: delta_W = A @ B * scaling + + Args: + 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, in_features, out_features] + """ + return torch.bmm(lora_A, lora_B) * scaling 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 6409b5b0..7f651f02 100644 --- a/src/xorl/qlora/modules/moe_experts.py +++ b/src/xorl/qlora/modules/moe_experts.py @@ -32,7 +32,7 @@ from transformers.utils import cached_file from xorl.lora.modules.base import LoraModule -from xorl.ops.group_gemm.kernel import compute_lora_scaling +from xorl.ops.group_gemm.kernel.lora_utils import compute_lora_scaling from xorl.ops.quantize import ( block_fp8_dequantize_gkn, block_fp8_quantize_gkn, @@ -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/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 7015ad7e..7915eadf 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( @@ -896,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 2d799c86..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.""" @@ -190,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: @@ -200,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, ) @@ -295,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, @@ -325,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, @@ -435,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, @@ -483,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 @@ -569,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. @@ -585,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( @@ -619,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: """ @@ -642,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}") @@ -719,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}") @@ -753,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. @@ -761,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") @@ -778,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. @@ -835,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 @@ -862,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)})") @@ -888,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 7482d4e7..fe15e31e 100644 --- a/src/xorl/server/api_server/orchestrator_client.py +++ b/src/xorl/server/api_server/orchestrator_client.py @@ -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 diff --git a/src/xorl/server/api_server/server.py b/src/xorl/server/api_server/server.py index 2ee23132..47fc30db 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,10 +160,11 @@ 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 @@ -171,10 +189,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 +200,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. @@ -269,7 +312,7 @@ def _build_info(result: Dict[str, Any]) -> Dict[str, Any]: return {"auto_loaded": True, "auto_load_path": result.get("auto_load_path")} return {} - 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 +322,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 +370,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 +382,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 +402,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 +508,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..85a9aea5 100644 --- a/src/xorl/server/api_server/training_ops.py +++ b/src/xorl/server/api_server/training_ops.py @@ -5,7 +5,7 @@ import logging import math import time -from typing import Any, Dict +from typing import Any, Dict, Optional from fastapi import HTTPException, status @@ -51,6 +51,25 @@ 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 + # ========================================================================= # Two-Phase Request Pattern Methods # ========================================================================= @@ -361,7 +380,7 @@ async def forward(self, request: ForwardRequest) -> 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 +406,39 @@ async def optim_step(self, request: OptimStepRequest) -> OptimStepResponse: self._require_engine() try: + adam_params = request.adam_params + lr = request.learning_rate + if lr is None and adam_params is not None: + lr = adam_params.learning_rate + if lr is None: + lr = self._session_default_learning_rate(request.model_id) + if lr is None: + lr = self._server_default_learning_rate() + if lr is None: + 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" + ), + ) + # 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 +474,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 e577a62f..f1edab96 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, ) @@ -358,6 +378,18 @@ async def sync_inference_weights( 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): return await self._execute( "register_adapter", diff --git a/src/xorl/server/launcher.py b/src/xorl/server/launcher.py index e56d85a1..010788bb 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 @@ -245,10 +246,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 +272,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 +314,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 +671,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 @@ -1034,6 +1070,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 +1092,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 +1125,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 f7379e2d..5d6552a3 100644 --- a/src/xorl/server/orchestrator/orchestrator.py +++ b/src/xorl/server/orchestrator/orchestrator.py @@ -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 @@ -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..567245f7 100644 --- a/src/xorl/server/orchestrator/packing.py +++ b/src/xorl/server/orchestrator/packing.py @@ -716,7 +716,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 +746,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,22 +768,46 @@ 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) diff --git a/src/xorl/server/orchestrator/request_processor.py b/src/xorl/server/orchestrator/request_processor.py index 14fb7e66..353ee8fa 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, @@ -195,10 +196,10 @@ async def _execute_model_pass( loss_fn=loss_fn, loss_fn_params=loss_fn_params, model_id=p.model_id, + routed_experts=p.routed_experts, + routed_expert_logits=p.routed_expert_logits, request_id=request.request_id, ) - if op_name == "forward_backward": - kwargs["routed_experts"] = p.routed_experts result = await backend_method(**kwargs) @@ -626,6 +627,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 5846fee3..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" diff --git a/src/xorl/server/protocol/operations.py b/src/xorl/server/protocol/operations.py index 395ecca2..309a5644 100644 --- a/src/xorl/server/protocol/operations.py +++ b/src/xorl/server/protocol/operations.py @@ -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..c0a339da 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,319 @@ 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) + 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": { + "type": self.optimizer_type, + "learning_rate": float(lr), + "weight_decay": float(self.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), + }, + } + + 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 +446,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 +475,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 +487,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 +520,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 +550,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 +609,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 +684,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 +734,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 +807,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 +916,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 +936,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 +949,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 +1032,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 acaabe5e..7aa170c8 100644 --- a/src/xorl/server/runner/checkpoint/manager.py +++ b/src/xorl/server/runner/checkpoint/manager.py @@ -32,6 +32,7 @@ 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 @@ -164,6 +165,59 @@ def _build_dcp_load_state(self, checkpoint_path: str, load_optimizer: bool) -> D # 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. @@ -185,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. @@ -202,11 +295,16 @@ def _save_lora_weights(self, save_path: str, model_id: str) -> None: self._adapter_manager.switch_adapter(model_id, auto_register=True) # Use fast adapter-manager path when available (avoids FSDP unshard). - # Skip the fast path for MoE LoRA: it reads rank-local params without an - # all_gather across EP ranks, so it would export only num_local_experts - # (e.g. 32) instead of num_experts (e.g. 128) for per-expert MoE tensors. - is_moe_lora = bool(self.lora_config.get("moe_hybrid_shared_lora", False)) or any( - m in (self.lora_config.get("lora_target_modules") or []) for m in ("gate_proj", "up_proj", "down_proj") + # 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") @@ -221,18 +319,33 @@ def _save_lora_weights(self, save_path: str, model_id: str) -> None: # 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 @@ -270,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, @@ -614,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: @@ -941,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: diff --git a/src/xorl/server/runner/model_runner.py b/src/xorl/server/runner/model_runner.py index c1cd9901..350ce34d 100644 --- a/src/xorl/server/runner/model_runner.py +++ b/src/xorl/server/runner/model_runner.py @@ -17,7 +17,9 @@ import logging import math import os +import shutil import time +from copy import deepcopy from typing import Any, Dict, List, Optional import torch @@ -44,6 +46,7 @@ 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, @@ -93,45 +96,72 @@ def configure_rank0_logging(logger_instance, rank): logger_instance.addFilter(RankFilter(rank)) +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], - metric_ops: Optional[Dict[str, str]], sp_group, + metric_ops: Optional[Dict[str, str]] = None, ) -> Dict[str, Any]: """ All-reduce KL/ratio metrics across the sequence-parallel (Ulysses) group. - With Ulysses SP, each rank sees only a shard of the sequence. Rank 0 often - has only prompt tokens (target_tokens=-100), so its KL stats are zeros. - Reducing here makes per-mb metrics CP-global before downstream cross-mb / - cross-DP accumulation. Means are raw partial sums; min/max use Β±inf - identity. Nothing is finalized β€” downstream sum-then-divide still applies. + With Ulysses SP, each rank only sees a shard of the sequence. Rank 0 often + has only prompt tokens (all target_tokens=-100), so its KL stats are zeros. + 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 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. """ - metric_ops = metric_ops or {} - if not metrics: - return metrics + # 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 - device = torch.device("cuda") + 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 - by_op: Dict[str, list[str]] = {"mean": [], "min": [], "max": []} - for k in metrics: - # valid_tokens is folded into the mean stack β€” it SUM-reduces too. - by_op.setdefault(metric_ops.get(k, "mean"), []).append(k) + 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() - reduced: Dict[str, torch.Tensor] = {} - for op_name, keys in by_op.items(): - if not keys: + for key, value in list(metrics.items()): + if key in {"valid_tokens", "_n_valid_kl"}: continue - stacked = torch.stack([torch.as_tensor(metrics[k], dtype=torch.float64, device=device) for k in keys]) - reduce_op = {"min": dist.ReduceOp.MIN, "max": dist.ReduceOp.MAX}.get(op_name, dist.ReduceOp.SUM) - dist.all_reduce(stacked, op=reduce_op, group=sp_group) - for i, k in enumerate(keys): - reduced[k] = stacked[i] - - metrics.update(reduced) - # valid_tokens is a Python int downstream (used as a denominator). - if "valid_tokens" in reduced: - metrics["valid_tokens"] = int(reduced["valid_tokens"].item()) + 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) + return metrics @@ -192,6 +222,7 @@ 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 " @@ -220,6 +251,8 @@ def __init__( # 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) @@ -235,7 +268,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 +285,10 @@ def __init__( 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,13 +298,26 @@ 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 @@ -306,6 +358,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: @@ -338,13 +465,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: @@ -354,10 +478,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: @@ -477,7 +641,7 @@ def _initialize_model(self): 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), @@ -503,6 +667,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 @@ -571,33 +736,91 @@ def _resolve_lora_target_modules(self) -> List[str]: 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") 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 = 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", - "muon_distributed_mode", - ) - if k in self.train_config - } + optimizer_kwargs = self._get_optimizer_kwargs() self.optimizer = build_optimizer( self.model, lr=self.train_config.get("lr", 1e-5), @@ -605,7 +828,7 @@ 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), ) @@ -657,16 +880,13 @@ def _load_initial_checkpoint(self) -> None: self._checkpoint_mgr.load_state(checkpoint_path, load_optimizer=True) self._sync_from_checkpoint_state() - def register_lora_adapter(self, model_id: str, lr: float) -> Dict[str, Any]: + def register_lora_adapter(self, model_id: str, lr: Optional[float]) -> Dict[str, Any]: """ - Register a new LoRA adapter for a training run. - - Creates fresh LoRA weights and optimizer state for the given model_id. - If the model_id already has an adapter, it will be replaced. + 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 @@ -677,11 +897,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, @@ -690,7 +921,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) @@ -754,82 +985,83 @@ 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, metric_ops=None): - """Accumulate IS metrics across micro-batches. + 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) - ``metric_ops`` tags non-mean keys; means accumulate by sum (finalized as - sum/count in _finalize_is_metrics), min/max by torch.minimum/maximum. - Values stay on-device β€” one collective + one ``.item()`` per metric in - _finalize_is_metrics. - """ + @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 = float(new_metrics.get("valid_tokens", 1)) - for k, v in new_metrics.items(): - op = metric_ops.get(k, "mean") - v_t = torch.as_tensor(v, dtype=torch.float64, device="cuda") - if op in ("min", "max"): - entry = accumulated.get(k) - if entry is None: - accumulated[k] = {"value": v_t.clone(), "op": op} + 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: + 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: - entry["value"] = ( - torch.minimum(entry["value"], v_t) if op == "min" else torch.maximum(entry["value"], v_t) - ) + 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: - entry = accumulated.get(k) - if entry is None: - accumulated[k] = { - "sum": v_t.clone(), - "count": 1.0 if k == "valid_tokens" else n_tokens, - "op": "mean", - } - else: - entry["sum"] = entry["sum"] + v_t - entry["count"] += 1.0 if k == "valid_tokens" else n_tokens + accumulated[k]["sum"] += v + accumulated[k]["count"] += n_tokens @staticmethod def _finalize_is_metrics(accumulated, result): - """All-reduce IS metrics across DP, then write finalized values to result. - - One SUM-allreduce for mean partial-sums concatenated with their counts; - one MIN/MAX-allreduce per non-mean group. Min/max with non-finite - reductions (every rank empty) fall back to 1.0. - """ + """All-reduce IS metrics across DP group, then add averaged values to result dict.""" if not accumulated: return ps = get_parallel_state() - dp_group = ps.dp_group if ps.dp_enabled else None - - groups: Dict[str, list[str]] = {"mean": [], "min": [], "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="cuda") - if dp_group is not None: - sums_and_counts = torch.cat([sums, counts]) - dist.all_reduce(sums_and_counts, op=dist.ReduceOp.SUM, group=dp_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[f"is_{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 dp_group is not None: - dist.all_reduce(stacked, op=reduce_op, group=dp_group) - values = stacked.tolist() - for k, v in zip(keys, values): - result[f"is_{k}"] = v if math.isfinite(v) else 1.0 + device = torch.device(get_device_type()) + if ps.dp_enabled: + dp_group = ps.dp_group + for k, v in accumulated.items(): + 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() + 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() + v["count"] = 1 if math.isfinite(v["sum"]) else 0 + else: + 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(): + 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): """Count valid tokens across all micro-batches and all-reduce across DP group. @@ -850,7 +1082,7 @@ def _count_active_microbatches(self, micro_batches) -> tuple[int, int]: # ========================================================================= def _compute_micro_batch_loss(self, micro_batch, loss_fn, loss_fn_params): - """Compute loss for a single micro-batch. Returns (local_loss_sum, per_token_outputs_dict, is_metrics, metric_ops, model_outputs).""" + """Compute loss for a single micro-batch.""" params = loss_fn_params or {} return_per_token = params.get("return_per_token", True) @@ -860,14 +1092,11 @@ def _compute_micro_batch_loss(self, micro_batch, loss_fn, loss_fn_params): outputs = self.model(**model_inputs, use_cache=False, output_hidden_states=False) 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)) + metric_sum_reducer = TokenPartial(scale=hidden_states.new_tensor(1.0, dtype=torch.float32)) per_token_outputs = {} is_metrics = None - metric_ops = None + is_metric_ops = None if loss_fn in ["causallm_loss", "cross_entropy"]: labels = micro_batch.get("labels") @@ -878,9 +1107,8 @@ def _compute_micro_batch_loss(self, micro_batch, loss_fn, loss_fn_params): return_per_token=return_per_token, ce_mode=self.ce_mode, lm_head_fp32=self.lm_head_fp32, - loss_reducer=token_sum_reducer, ) - local_loss_sum = _result.loss + loss = _result.loss if return_per_token: per_token_outputs["logprobs"] = _result.per_token_logprobs per_token_outputs["loss"] = _result.per_token_loss @@ -900,16 +1128,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, - loss_reducer=token_sum_reducer, - metric_reducer=token_sum_reducer, + metric_reducer=metric_sum_reducer, ) - local_loss_sum = _result.loss + loss = _result.loss per_token_outputs["logprobs"] = _result.per_token_logprobs is_metrics = _result.metrics - metric_ops = _result.metric_ops + 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, metric_ops, 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) @@ -1001,21 +1234,24 @@ 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, - loss_reducer=token_sum_reducer, - metric_reducer=token_sum_reducer, + metric_reducer=metric_sum_reducer, ) - local_loss_sum = _result.loss + loss = _result.loss per_token_outputs["logprobs"] = _result.per_token_logprobs is_metrics = _result.metrics - metric_ops = _result.metric_ops + 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, metric_ops, get_parallel_state().ulysses_group) + is_metrics = _sp_allreduce_kl_metrics( + is_metrics, + get_parallel_state().ulysses_group, + is_metric_ops, + ) else: raise ValueError(f"Unknown loss_fn: {loss_fn}") - return local_loss_sum, per_token_outputs, is_metrics, metric_ops, outputs + return loss, per_token_outputs, is_metrics, is_metric_ops, outputs # ========================================================================= # Per-sample K3 KL divergence @@ -1084,10 +1320,11 @@ 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) # Count valid tokens globally global_valid_tokens = self._count_global_valid_tokens(micro_batches) - if compute_backward and self._use_distsignsgd: + 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 @@ -1123,16 +1360,16 @@ def _forward_loop( # Forward pass + loss computation with self.model_fwd_context: - local_loss_sum, per_token_outputs, is_metrics, metric_ops, outputs = self._compute_micro_batch_loss( + loss, per_token_outputs, is_metrics, is_metric_ops, outputs = self._compute_micro_batch_loss( micro_batch, loss_fn, params ) logger.debug( f"Rank {self.rank}: micro_batch {batch_idx}/{len(micro_batches)} " - f"local_loss_sum={local_loss_sum.item():.6f}, local_valid_tokens={local_valid_tokens.item()}, " + f"loss={loss.item():.6f}, local_valid_tokens={local_valid_tokens.item()}, " f"global_valid_tokens={global_valid_tokens.item()}" ) - # Note: local_loss_sum is always finite even when local_valid_tokens=0, because + # Note: loss is always finite even when local_valid_tokens=0, because # causallm_loss_function uses reduction="none" + manual mean with # clamp(min=1) denominator. No need to replace with zeros_like # (which would break the autograd graph and cause FSDP2 deadlocks). @@ -1178,12 +1415,16 @@ def _forward_loop( } ) - # Backward + reporting on the raw partial sum: cross-mb / cross-DP - # accumulation composes under SUM-allreduce, then optim_step divides - # once by global_valid_tokens. FSDP grad averaging is off - # (set_gradient_divide_factor(1.0) in torch_parallelize). + # Gradient accumulation β€” raw (unnormalized) backward. + # Normalization by total accumulated valid tokens is deferred to optim_step. + # FSDP's automatic gradient averaging is disabled (set_gradient_divide_factor(1.0) + # in torch_parallelize), so no fsdp_size compensation is needed here. + # When local_valid_tokens=0, this produces 0 gradients while preserving + # the full autograd graph through all parameters (including lm_head weight), + # 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") @@ -1193,21 +1434,26 @@ def _forward_loop( set_replay_stage("replay_backward") with self.model_bwd_context: - local_loss_sum.backward() + raw_loss.backward() + # Loss reporting (separately, no grad): compute normalized per-token loss with torch.no_grad(): - loss_report = local_loss_sum.detach() + loss_report = loss.detach() * local_valid_tokens 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 += local_loss_sum.item() / global_valid_tokens.item() + total_loss += loss.item() * (local_valid_tokens.item() / global_valid_tokens.item()) - self._accumulate_is_metrics(accumulated_is_metrics, is_metrics, metric_ops) + # Accumulate IS metrics + self._accumulate_is_metrics(accumulated_is_metrics, is_metrics, is_metric_ops) # Cleanup - del micro_batch, outputs, local_loss_sum + del micro_batch, outputs, loss + if compute_backward: + del raw_loss # Note: gc.collect() + empty_cache() removed from per-step path. # They cost ~250ms + ~50ms per step (profiled on Qwen3-8B 8xH100). @@ -1222,13 +1468,13 @@ def _forward_loop( sync_sp_gradients( self.model, get_parallel_state().sp_grad_sync_group, - skip_dtensor_grads=self._use_distsignsgd, + 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 self._use_distsignsgd: + if use_distsignsgd: self._accumulated_active_microbatches[model_id] = ( self._accumulated_active_microbatches.get(model_id, 0) + active_microbatches ) @@ -1380,7 +1626,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 @@ -1390,6 +1636,7 @@ 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) # 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 @@ -1480,7 +1727,7 @@ def forward_backward( # PP path if self.pp_enabled: global_valid_tokens = self._count_global_valid_tokens(micro_batches) - if self._use_distsignsgd: + if use_distsignsgd: active_microbatches, active_voter_total = self._count_active_microbatches(micro_batches) else: active_microbatches, active_voter_total = 0, 0 @@ -1504,7 +1751,7 @@ 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 self._use_distsignsgd: + if use_distsignsgd: self._accumulated_active_microbatches[model_id] = ( self._accumulated_active_microbatches.get(model_id, 0) + active_microbatches ) @@ -1587,11 +1834,19 @@ def forward( 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() @@ -1604,20 +1859,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 @@ -1668,8 +1930,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 = self._accumulated_active_microbatches.pop(model_id, 0) - accumulated_active_voter_total = self._accumulated_active_voter_total.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: @@ -1688,12 +1951,13 @@ 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: - if self._use_distsignsgd: + if use_distsignsgd: if accumulated_active_voter_total > 0: scale_model_gradients( self.model, @@ -1711,7 +1975,7 @@ def optim_step( ps = get_parallel_state() clip_value = get_effective_grad_clip_value( clip_value, - use_distsignsgd=self._use_distsignsgd, + use_distsignsgd=use_distsignsgd, ) grad_norm = clip_gradients( @@ -1768,6 +2032,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), @@ -1801,6 +2070,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): @@ -1840,6 +2110,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 @@ -1882,6 +2155,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..0d62b2ce 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( @@ -890,6 +891,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 +1014,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 +1073,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 +1137,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 +1286,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/server_arguments.py b/src/xorl/server/server_arguments.py index cbce8994..91a78b33 100644 --- a/src/xorl/server/server_arguments.py +++ b/src/xorl/server/server_arguments.py @@ -11,6 +11,7 @@ from dataclasses import dataclass, field from typing import Any, Dict, List, Literal, Optional +import torch import yaml from xorl.ops.loss import CrossEntropyMode @@ -143,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"}) @@ -206,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"}) @@ -214,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( @@ -257,6 +275,16 @@ class ServerArguments: }, ) + 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."}, @@ -478,6 +506,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( @@ -523,6 +559,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 @@ -544,6 +587,36 @@ class ServerArguments: }, ) + @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 @@ -567,6 +640,25 @@ 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." @@ -610,10 +702,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, @@ -626,6 +720,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, @@ -634,6 +729,8 @@ 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, @@ -650,6 +747,7 @@ def to_config_dict(self) -> Dict[str, Any]: "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, @@ -669,6 +767,7 @@ 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, @@ -679,6 +778,7 @@ def to_config_dict(self) -> Dict[str, Any]: "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..0caf2d1c --- /dev/null +++ b/src/xorl/server/session_spec.py @@ -0,0 +1,442 @@ +"""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) + + 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 b724ecd0..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 (connect, send) +- SyncPushChannel: Sync PUSH socket (bind/connect, send) - SyncDealerChannel: Sync DEALER socket (connect, poll, recv) -- AsyncPullChannel: Async PULL socket (bind, 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) """ diff --git a/src/xorl/server/weight_sync/handler.py b/src/xorl/server/weight_sync/handler.py index 572dbbee..49063626 100644 --- a/src/xorl/server/weight_sync/handler.py +++ b/src/xorl/server/weight_sync/handler.py @@ -403,6 +403,45 @@ def _add_rank_timing_breakdown( timing_breakdown["min_rank_transfer_s"] = min_transfer_s timing_breakdown["rank_transfer_spread_s"] = max_transfer_s - min_transfer_s + @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 # ======================================================================== @@ -2004,9 +2043,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 @@ -2149,9 +2186,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 @@ -2302,9 +2337,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" diff --git a/src/xorl/trainers/model_builder.py b/src/xorl/trainers/model_builder.py index 362e08e3..bcdff8a1 100644 --- a/src/xorl/trainers/model_builder.py +++ b/src/xorl/trainers/model_builder.py @@ -149,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. @@ -190,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, ) diff --git a/src/xorl/trainers/trainer.py b/src/xorl/trainers/trainer.py index 7f40cb37..05d4579a 100644 --- a/src/xorl/trainers/trainer.py +++ b/src/xorl/trainers/trainer.py @@ -446,6 +446,7 @@ def _build_model(self) -> None: 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 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/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..f8177fd7 100644 --- a/tests/e2e/server_utils.py +++ b/tests/e2e/server_utils.py @@ -369,14 +369,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/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/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_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/server/api_server/test_api_server.py b/tests/server/api_server/test_api_server.py index a3d46897..d7be18ae 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") @@ -178,6 +237,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 4d7d0707..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", 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/orchestrator/test_packing.py b/tests/server/orchestrator/test_packing.py index 7fe6c9c2..fb93e40f 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 # ============================================================================ @@ -77,7 +77,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 +86,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..ffdc8de2 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,10 @@ LoadStateData, ModelPassData, OptimStepData, + RegisterSessionData, SaveStateData, ) +from xorl.server.runner.runner_dispatcher import RunnerDispatcher # ============================================================================ @@ -108,6 +112,167 @@ async def test_forward_backward_operations(processor): assert "loss" in output.outputs[0] +@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 +331,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_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_lora_checkpoint_roundtrip.py b/tests/server/runner/test_lora_checkpoint_roundtrip.py index 8cb7b550..e331c397 100644 --- a/tests/server/runner/test_lora_checkpoint_roundtrip.py +++ b/tests/server/runner/test_lora_checkpoint_roundtrip.py @@ -8,7 +8,13 @@ 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 @@ -27,35 +33,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): @@ -82,6 +88,74 @@ 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) 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_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 7ca35b29..3d2f9b7b 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( @@ -55,6 +106,52 @@ def test_load_server_arguments_threads_distsignsgd_through_nested_config(tmp_pat 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( @@ -76,6 +173,51 @@ def test_load_server_arguments_rejects_broadcast_load_weights_mode(tmp_path): 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): config_path = tmp_path / "server_config.yaml" config_path.write_text( @@ -86,6 +228,8 @@ 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, @@ -101,8 +245,18 @@ 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" @@ -115,3 +269,40 @@ def test_load_server_arguments_threads_muon_gram_newton_schulz_through_nested_co 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_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 From 31fc4e81106f825c21e104139fc256df41827739 Mon Sep 17 00:00:00 2001 From: Qingyang Wu Date: Mon, 18 May 2026 22:01:34 -0700 Subject: [PATCH 35/49] perf(runner): use torch.inference_mode() for ModelRunner.forward eval path MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit inference_mode is strictly stronger than no_grad: it additionally skips view tracking and version counters. The eval-only forward path's returned tensors are only consumed by logging and cross-rank reductions that never re-enter autograd, so the strictness is safe here. Sleep/wake_up paths intentionally left on no_grad β€” they manipulate parameters that subsequently re-enter autograd-tracked code. Refs. --- src/xorl/server/runner/model_runner.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/xorl/server/runner/model_runner.py b/src/xorl/server/runner/model_runner.py index 350ce34d..b03385fa 100644 --- a/src/xorl/server/runner/model_runner.py +++ b/src/xorl/server/runner/model_runner.py @@ -1828,7 +1828,10 @@ def forward_backward( return result - @torch.no_grad() + # inference_mode (vs no_grad) skips view tracking and version counters, + # which is safe here because the returned tensors are only consumed by + # logging / cross-rank reductions that never re-enter autograd. + @torch.inference_mode() def forward( self, micro_batches: List[Dict[str, Any]], From 9bb11383e049a25a646b0a792f302d23f2a30c70 Mon Sep 17 00:00:00 2001 From: Ashwinee Panda Date: Tue, 19 May 2026 01:35:29 -0700 Subject: [PATCH 36/49] fix(server): align R3 routing replay with SP microbatches * Fix R3 routing replay SP alignment * Align R3 routing replay with ring attention --------- Co-authored-by: Qingyang Wu --- .../runner/utils/routing_replay_handler.py | 226 ++++++++++++++---- .../runner/test_routing_replay_handler.py | 139 +++++++++++ 2 files changed, 318 insertions(+), 47 deletions(-) create mode 100644 tests/server/runner/test_routing_replay_handler.py diff --git a/src/xorl/server/runner/utils/routing_replay_handler.py b/src/xorl/server/runner/utils/routing_replay_handler.py index 407d95fb..1e13c12b 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 @@ -323,11 +322,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 +453,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/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] From c5ab17ec140fd7735e081a7750dd4ac28b25b9d0 Mon Sep 17 00:00:00 2001 From: Ashwinee Panda Date: Wed, 20 May 2026 14:33:57 -0700 Subject: [PATCH 37/49] perf(weight-sync): improve P2P sync * Improve P2P weight sync efficiency * Document P2P weight sync env profile * chore(weight-sync): apply lint formatting * chore: retrigger PR checks * test(weight-sync): update async transfer helper calls * fix(weight-sync): cover qwen fp8 p2p sync gaps * docs(weight-sync): clarify fp8 p2p setup --- src/xorl/server/weight_sync/README.md | 115 ++- src/xorl/server/weight_sync/backends/p2p.py | 898 +++++++++++++++--- src/xorl/server/weight_sync/handler.py | 455 ++++++++- .../weight_sync/test_fp8_quantization.py | 15 + .../server/weight_sync/test_handler_config.py | 286 +++++- .../server/weight_sync/test_p2p_async_api.py | 8 + .../weight_sync/test_p2p_backend_protocol.py | 754 ++++++++++++++- 7 files changed, 2288 insertions(+), 243 deletions(-) diff --git a/src/xorl/server/weight_sync/README.md b/src/xorl/server/weight_sync/README.md index 2d7aee47..a5441154 100644 --- a/src/xorl/server/weight_sync/README.md +++ b/src/xorl/server/weight_sync/README.md @@ -241,11 +241,71 @@ 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: -- `P2P_SYNC_QUANTIZATION='{"quant_method":"fp8","fmt":"e4m3","weight_block_size":[128,128]}'`: - quantizes projection weights on the trainer side and transfers FP8 weights - plus `weight_scale_inv` tensors to an FP8 SGLang receiver. +- 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. @@ -270,18 +330,59 @@ P2P tuning options: 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_WEIGHT_SYNC_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_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 staged P2P bucket; the default is 4 GiB. + 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 diff --git a/src/xorl/server/weight_sync/backends/p2p.py b/src/xorl/server/weight_sync/backends/p2p.py index a5f04ac4..ddb171a0 100644 --- a/src/xorl/server/weight_sync/backends/p2p.py +++ b/src/xorl/server/weight_sync/backends/p2p.py @@ -17,11 +17,13 @@ """ import dataclasses +import ipaddress import logging import os import socket import time -from concurrent.futures import Future, ThreadPoolExecutor +import zlib +from concurrent.futures import Future, ThreadPoolExecutor, as_completed from typing import Any, Dict, FrozenSet, List, Optional, Tuple import requests @@ -64,7 +66,26 @@ class _StagedTransfer: peer_ptr: int nbytes: int memory_handle: Optional[int] - debug_entries: List[_TransferDebugEntry] + 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 @@ -113,6 +134,7 @@ def throughput_mb_s(self) -> float: _HTTP_TIMEOUT_SECONDS = 600 +_TRANSFER_DEBUG_SAMPLE_LIMIT = 6 def _env_float(name: str, default: float) -> float: @@ -145,6 +167,19 @@ def _env_int(name: str, default: int, *, minimum: int = 1) -> int: 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) @@ -153,6 +188,31 @@ 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 @@ -171,7 +231,9 @@ def _align_up(value: int, alignment: int) -> int: # 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). +# 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))) @@ -274,25 +336,52 @@ def _transfer_debug_entry(name: str, loc: Dict[str, Any], nbytes: int) -> _Trans ) -def _format_transfer_debug(debug_entries: List[_TransferDebugEntry]) -> str: +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[:6]: + 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 len(debug_entries) > 6: - parts.append(f"... {len(debug_entries) - 6} more") + 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], List[List[_TransferDebugEntry]]]], + by_session: Dict[str, Tuple[List[int], List[int], List[int], Optional[List[_TransferDebugSample]]]], session_id: str, i: int, end: int, @@ -300,20 +389,26 @@ def _chunk_sizes( return by_session[session_id][2][i:end] -def _chunk_debug_entries( - by_session: Dict[str, Tuple[List[int], List[int], List[int], List[List[_TransferDebugEntry]]]], +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, -) -> List[_TransferDebugEntry]: - debug_lists = by_session[session_id][3][i:end] - return [entry for entries in debug_lists for entry in entries] +) -> 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], List[List[_TransferDebugEntry]]]], + 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], @@ -335,7 +430,7 @@ def _run_sync_transfer_items( 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_entries(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) @@ -344,7 +439,7 @@ def _run_sync_transfer_items( def _run_async_transfer_items( *, engine_wrapper: Any, - by_session: Dict[str, Tuple[List[int], List[int], List[int], List[List[_TransferDebugEntry]]]], + 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], @@ -377,7 +472,7 @@ def _run_async_transfer_items( 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_entries(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: @@ -396,7 +491,7 @@ def _run_async_transfer_items( 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_entries(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: @@ -416,7 +511,7 @@ def _run_async_transfer_items( 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_entries(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: @@ -431,7 +526,7 @@ def _run_async_transfer_items( def _transfer_small_entries( *, engine_wrapper: Any, - small_session_data: Dict[str, List[Tuple[int, int, int, _TransferDebugEntry]]], + 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], @@ -447,19 +542,35 @@ def _transfer_small_entries( 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 src_ptr, peer_ptr, nbytes, debug_entry in triples: - total_bytes += nbytes - session_bytes[session_id] = session_bytes.get(session_id, 0) + nbytes - num_buffers += 1 - ret = engine_wrapper.batch_transfer_sync(session_id, [src_ptr], [peer_ptr], [nbytes]) - if ret < 0: + 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={ret} (nbytes={nbytes}, bucket {bucket_idx}, " - f"{_format_transfer_debug([debug_entry])}, " - f"session_info={session_debug_info.get(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: @@ -476,16 +587,18 @@ def _do_async_transfer( *, engine_wrapper: Any, copy_done_event: "torch.cuda.Event", - by_session: Dict[str, Tuple[List[int], List[int], List[int], List[List[_TransferDebugEntry]]]], - small_session_data: Dict[str, List[Tuple[int, int, int, _TransferDebugEntry]]], + 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. @@ -524,7 +637,7 @@ def _do_async_transfer( # 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 os.environ.get("XORL_P2P_USE_ASYNC_API", "0") == "1": + if use_async_api: _run_sync_transfer_items( engine_wrapper=engine_wrapper, by_session=by_session, @@ -571,15 +684,16 @@ def _do_async_transfer( timing.num_small_buffers = num_small_buffers timing.session_bytes = session_bytes timing.session_transfer_s = session_transfer_s - 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, - ) + 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): @@ -623,14 +737,15 @@ def __init__(self, config: TransportConfig, **kwargs: Any) -> None: # registered with Mooncake once at first use and reused for # every subsequent bucket. # - # Default 2 pools (ping-pong). Bumped via XORL_P2P_NUM_POOLS; - # combined with XORL_P2P_MOONCAKE_WORKERS=N this gives N-way - # concurrent Mooncake calls per rank, hiding per-call latency - # on medium-speed nodes. + # 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 @@ -643,9 +758,13 @@ def __init__(self, config: TransportConfig, **kwargs: Any) -> None: # 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)) @@ -656,6 +775,25 @@ def __init__(self, config: TransportConfig, **kwargs: Any) -> None: # 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 @@ -673,6 +811,7 @@ def __init__(self, config: TransportConfig, **kwargs: Any) -> None: # 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 @@ -724,10 +863,248 @@ def _session_ids_for_endpoint( 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 @@ -749,8 +1126,40 @@ def adopt_prepared_state( 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 = { - sid: {"session_id": sid, "adopted_from_rank0": True} for sid in self._receiver_session_ids + str(info["session_id"]): dict(info) for info in receiver_session_infos if info.get("session_id") } return True @@ -772,28 +1181,56 @@ def _initialize_multi_sender(self) -> bool: 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] * self._world_size + cached_states: List[bool] = [False] * group_world_size try: - dist.all_gather_object(cached_states, has_cached_state) + 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] * self._world_size + 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: - payload[0] = ("tensor_map", self._tensor_map, list(self._receiver_session_ids)) + 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 the broadcast. payload[0] is None on - # failure so non-zero ranks can short-circuit cleanly. - dist.broadcast_object_list(payload, src=0) + # 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: @@ -815,13 +1252,22 @@ def _initialize_multi_sender(self) -> bool: _, 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] * self._world_size + init_results: List[bool] = [False] * group_world_size try: - dist.all_gather_object(init_results, bool(local_ok)) + 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 @@ -848,7 +1294,7 @@ def _initialize_single_sender(self) -> bool: sender_session_id = self._engine.get_session_id() sender_info = { "session_id": sender_session_id, - "hostname": self._hostname, + "hostname": self._resolved_hostname or self._hostname, "gpu_id": self._gpu_id, "ib_device": self._engine.get_ib_device(), "training_rank": cfg.training_rank, @@ -857,64 +1303,101 @@ def _initialize_single_sender(self) -> bool: # 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. + # 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: - merged_tensor_map: Dict[str, List[Dict[str, Any]]] = { - name: [dict(loc) for loc in locators] for name, locators in self._tensor_map.items() - } + # 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 in self._receiver_session_ids: + 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 - for ep_idx, ep in enumerate(cfg.endpoints): - 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 request_cached_prepare: - payload["p2p_return_tensor_map"] = False - try: - resp = requests.post(url, json=payload, timeout=_prepare_timeout_seconds()) - except requests.RequestException as e: - logger.error(f"[P2P] /prepare_weights_update to {ep.host}:{ep.port} failed: {e}") - return False - if request_cached_prepare and resp.status_code in (400, 422): - logger.warning( - f"[P2P] cached prepare was rejected by {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 resp.status_code != 200: - logger.error( - f"[P2P] /prepare_weights_update returned {resp.status_code} " - f"from {ep.host}:{ep.port}: {resp.text}" - ) - return False - body = resp.json() - if not body.get("success", False): - logger.error(f"[P2P] prepare failed at {ep.host}:{ep.port}: {body.get('message')}") - return 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) @@ -937,6 +1420,7 @@ def _initialize_single_sender(self) -> bool: 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 @@ -956,6 +1440,7 @@ def _initialize_single_sender(self) -> bool: 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 @@ -979,7 +1464,9 @@ def _initialize_single_sender(self) -> bool: 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"(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 @@ -1087,6 +1574,7 @@ def complete_sync(self) -> None: 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" @@ -1098,6 +1586,8 @@ def complete_sync(self) -> None: } 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: @@ -1293,8 +1783,13 @@ def transfer_bucket( # pre-registered CPU pool. pending_small: Dict[str, List[_PendingTransfer]] = {} for sid, entries in list(pending.items()): - small = [e for e in entries if e.src_view.is_cuda and e.nbytes < self._cpu_pool_min_bytes] - large = [e for e in entries if (not e.src_view.is_cuda) or e.nbytes >= self._cpu_pool_min_bytes] + 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: @@ -1332,8 +1827,11 @@ def transfer_bucket( 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], List[List[_TransferDebugEntry]]]] = {} + 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() @@ -1351,35 +1849,45 @@ def transfer_bucket( for e in entries: if pool is None: raise RuntimeError("[P2P] CPU scratch pool was not initialized") - 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 + 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), - debug_entries=[_transfer_debug_entry(e.name, e.loc, e.nbytes)], + name=e.name, + loc=e.loc, ) ) - scratch_offset_bytes += e.nbytes # Coalesce: walk staged in sorted-by-peer order, merging # adjacent entries whose source pool ptr AND peer ptr are @@ -1390,7 +1898,7 @@ def transfer_bucket( src_ptrs: List[int] = [] peer_ptrs: List[int] = [] lens: List[int] = [] - debug_entries: List[List[_TransferDebugEntry]] = [] + debug_entries: Optional[List[_TransferDebugSample]] = [] if self._collect_transfer_debug else None memory_handles: List[Optional[int]] = [] for staged_entry in staged: if ( @@ -1401,48 +1909,76 @@ def transfer_bucket( and memory_handles[-1] == staged_entry.memory_handle ): lens[-1] += staged_entry.nbytes - debug_entries[-1].extend(staged_entry.debug_entries) + 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) - debug_entries.append(list(staged_entry.debug_entries)) + 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 total_post_coalesce < total_pre_coalesce: + 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, _TransferDebugEntry]]] = {} + 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_intervals: List[Tuple[int, int]] = [] + 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, _TransferDebugEntry]] = [] + 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()) - small_intervals.append((s_start, s_end)) - triples.append( - (int(sv.data_ptr()), e.peer_ptr, e.nbytes, _transfer_debug_entry(e.name, e.loc, e.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 - small_segs = self._intervals_per_cuda_segment(small_intervals) - small_register_ptrs = [iv[0] for iv in small_segs] - small_register_lens = [iv[1] - iv[0] for iv in small_segs] + 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 @@ -1462,8 +1998,6 @@ def transfer_bucket( else: copy_done_event = _CompletedCudaEvent() - chunk = max(1, int(os.environ.get("XORL_P2P_MOONCAKE_TRANSFER_CHUNK", "1"))) - if self._transfer_executor is None: raise RuntimeError("[P2P] transfer executor was not initialized") t_submit = time.perf_counter() @@ -1476,11 +2010,13 @@ def transfer_bucket( session_debug_info=self._session_debug_info, small_register_ptrs=small_register_ptrs, small_register_lens=small_register_lens, - chunk=chunk, + 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 @@ -1607,9 +2143,23 @@ def _percentile(values: List[float], percentile: float) -> float: @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 @@ -1709,7 +2259,7 @@ def merge(intervals: List[Tuple[int, int]]) -> List[Tuple[int, int]]: return merged cand_merged = merge(candidates) - registered = self._registered_intervals + registered = merge(self._registered_intervals) new: List[Tuple[int, int]] = [] for s, e in cand_merged: @@ -1730,6 +2280,40 @@ def merge(intervals: List[Tuple[int, int]]) -> List[Tuple[int, int]]: 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."): @@ -1891,6 +2475,7 @@ def _make_local_engine(self): # 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, @@ -1934,12 +2519,37 @@ def _resolve_local_hostname() -> str: Mooncake's handshake binds on this hostname; it must be reachable from the SGLang receiver. """ - # Prefer an FQDN that the inference side can route to. Fall back to the - # first non-loopback IPv4 address. + explicit = os.environ.get("XORL_P2P_HOSTNAME") + if explicit and explicit.strip(): + return explicit.strip() + + 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: - host = socket.getfqdn() - if host and host != "localhost": - return host + fqdn = socket.getfqdn() + if ip := _routable_ipv4(socket.gethostbyname(fqdn)): + return ip except Exception: pass - return socket.gethostbyname(socket.gethostname()) + return socket.gethostname() diff --git a/src/xorl/server/weight_sync/handler.py b/src/xorl/server/weight_sync/handler.py index 49063626..6e2cc9b3 100644 --- a/src/xorl/server/weight_sync/handler.py +++ b/src/xorl/server/weight_sync/handler.py @@ -31,7 +31,7 @@ import os import time from concurrent.futures import Future, ThreadPoolExecutor -from typing import Any, Dict, List, Optional, Tuple +from typing import Any, Callable, Dict, List, Optional, Tuple import torch import torch.distributed as dist @@ -77,8 +77,10 @@ def _env_int(name: str, default: int, *, minimum: int = 1) -> int: def _moe_bucket_size_bytes(sync_method: str) -> int: - """Default MoE bucket sizing is backend-specific; the env var remains an explicit override.""" + """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) @@ -209,6 +211,8 @@ def _safe_abort_token(value: Optional[str]) -> str: # --------------------------------------------------------------------------- _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 def _atexit_destroy_cached_backend() -> None: @@ -225,6 +229,107 @@ def _atexit_destroy_cached_backend() -> 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.""" @@ -241,6 +346,7 @@ def __init__(self, rank: int, world_size: int, trainer) -> None: 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() @@ -403,6 +509,29 @@ def _add_rank_timing_breakdown( 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, @@ -572,6 +701,8 @@ def _sync_weights( 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) @@ -591,6 +722,25 @@ def _sync_weights( # 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=[ @@ -632,9 +782,11 @@ def _sync_weights( self.rank, tuple( sorted( - (k, v) + (k, cache_v) for k, v in (_backend_config or {}).items() - if k not in {"flush_cache", "weight_version"} and isinstance(v, (str, int, bool, float)) + if k not in {"flush_cache", "weight_version", "process_group"} + for cache_v in (_backend_cache_value(v),) + if cache_v is not None ) ), ) @@ -718,12 +870,15 @@ def _add_rank_phase(name: str, start: float) -> None: # 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. - # Default off β€” flip via XORL_WEIGHT_SYNC_BATCH_MOE=1. + # 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 @@ -735,6 +890,21 @@ def _add_rank_phase(name: str, start: float) -> None: # 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 @@ -844,6 +1014,7 @@ def _add_rank_phase(name: str, start: float) -> None: 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) @@ -862,17 +1033,34 @@ def _add_rank_phase(name: str, start: float) -> None: 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) @@ -890,36 +1078,52 @@ def _add_rank_phase(name: str, start: float) -> None: # Stage 0: sender rank(s) broadcast directly to SGLang if _is_sender and current_buffer: t_phase = time.perf_counter() - current_buffer = self._unfuse_for_inference( - current_buffer, - model, - ) - if quantization and quantization.get("quant_method") == "fp8": - current_buffer = self._quantize_buffer_for_fp8( + if current_buffer: + current_buffer = self._unfuse_for_inference( current_buffer, - quantization_config=quantization, - target_device=self._fp8_quantization_target_device(backend), - phase_s=rank_phase_s, - phase_prefix="dense_fp8", + model, + ) + 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() - logger.info(f"Rank 0: [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() + 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 @@ -928,9 +1132,7 @@ def _add_rank_phase(name: str, start: float) -> None: # 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 and self.rank in backend.sender_ranks - ) + use_direct_ep = backend.supports_direct_ep_transfer for ctx in moe_contexts + ep_moe_contexts: if use_direct_ep: # batch_moe defers the per-call @@ -1160,6 +1362,8 @@ def _add_rank_phase(name: str, start: float) -> None: # 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. @@ -1376,13 +1580,16 @@ def _collect_ep_moe_data( 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. @@ -1422,6 +1629,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 = {} @@ -1532,19 +1750,24 @@ def _direct_ep_transfer_experts( 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"] - 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 - logger.info( f"Rank {self.rank}: [Direct-EP] prefix={full_prefix}, E_local={E_local}, E_total={E_local * ep_size}" ) @@ -1553,6 +1776,22 @@ def _direct_ep_transfer_experts( 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: @@ -1599,6 +1838,7 @@ def _direct_ep_transfer_experts( 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 @@ -1656,6 +1896,7 @@ def _direct_ep_transfer_experts( 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 @@ -1684,6 +1925,7 @@ def _direct_ep_transfer_experts( 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 @@ -2382,6 +2624,10 @@ def _fp8_cpu_workspace_min_capacity() -> int: 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") @@ -2439,7 +2685,15 @@ def _should_quantize_fp8_weight( name == prefix + ".weight" or name.startswith(prefix + ".") for prefix in modules_to_not_convert ) - return "_proj.weight" in name or name.endswith("fused_qkv_a_proj_with_mqa.weight") + 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: @@ -3146,12 +3400,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. @@ -3217,6 +3506,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" @@ -3286,10 +3586,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 @@ -3421,13 +3738,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/tests/server/weight_sync/test_fp8_quantization.py b/tests/server/weight_sync/test_fp8_quantization.py index 6f469a59..7beca75c 100644 --- a/tests/server/weight_sync/test_fp8_quantization.py +++ b/tests/server/weight_sync/test_fp8_quantization.py @@ -58,6 +58,21 @@ def test_fp8_quantization_includes_fused_mla_weight_by_default(): 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) diff --git a/tests/server/weight_sync/test_handler_config.py b/tests/server/weight_sync/test_handler_config.py index 2add0946..3dadeeb4 100644 --- a/tests/server/weight_sync/test_handler_config.py +++ b/tests/server/weight_sync/test_handler_config.py @@ -7,23 +7,307 @@ _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_moe_bucket_env_override_is_explicit(monkeypatch): +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, diff --git a/tests/server/weight_sync/test_p2p_async_api.py b/tests/server/weight_sync/test_p2p_async_api.py index 69a7f499..9e411f94 100644 --- a/tests/server/weight_sync/test_p2p_async_api.py +++ b/tests/server/weight_sync/test_p2p_async_api.py @@ -85,10 +85,12 @@ def test_async_api_success_status_zero_completes(monkeypatch): 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])] @@ -110,10 +112,12 @@ def test_async_api_uses_sync_fallback_for_medium_chunks(monkeypatch): 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 == [] @@ -136,10 +140,12 @@ def test_async_api_min_bytes_env_controls_cutoff(monkeypatch): 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])] @@ -162,10 +168,12 @@ def test_async_api_status_poll_timeout(monkeypatch): 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, ) diff --git a/tests/server/weight_sync/test_p2p_backend_protocol.py b/tests/server/weight_sync/test_p2p_backend_protocol.py index db19ab25..314feeb1 100644 --- a/tests/server/weight_sync/test_p2p_backend_protocol.py +++ b/tests/server/weight_sync/test_p2p_backend_protocol.py @@ -18,6 +18,7 @@ from __future__ import annotations +import socket import sys import types from concurrent.futures import Future @@ -28,7 +29,12 @@ import torch from xorl.server.weight_sync.backends.base import EndpointConfig, TransportConfig -from xorl.server.weight_sync.backends.p2p import P2PTransportBackend +from xorl.server.weight_sync.backends.p2p import ( + P2PTransportBackend, + _async_api_enabled, + _resolve_local_hostname, + _transfer_small_entries, +) class _FakeResponse: @@ -85,6 +91,69 @@ def batch_transfer_sync( 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)], @@ -173,7 +242,8 @@ def test_prepare_payload_uses_p2p_transport_and_engine_info(self): assert body["group_name"] == "weight_sync_group" assert "p2p_return_tensor_map" not in body - def test_initialize_aggregates_tensor_map_across_endpoints(self): + 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( @@ -282,6 +352,7 @@ def test_cached_prepare_can_reuse_existing_tensor_map(self): 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 @@ -349,7 +420,8 @@ def test_cached_prepare_retries_full_map_when_receiver_session_changes(self): 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): + 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 = { @@ -462,7 +534,8 @@ def test_cached_prepare_retries_full_map_when_receiver_rejects_flag(self): 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_restarts_all_endpoints_when_later_endpoint_rejects_flag(self): + 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 = { @@ -495,17 +568,6 @@ def test_cached_prepare_restarts_all_endpoints_when_later_endpoint_rejects_flag( } backend._receiver_session_ids = ["recv-0:7000", "recv-1:7000"] backend._prefer_cached_prepare = True - ep0_full_locator = { - **_hf_locator( - tp_rank=0, - full_shape=[128, 64], - slc=[[0, 64], [0, 64]], - ptr=0x1000, - nbytes=64 * 64 * 2, - session_id="recv-0:7000", - ), - "hf_name": param_name, - } ep1_full_locator = { **_hf_locator( tp_rank=1, @@ -526,15 +588,6 @@ def test_cached_prepare_restarts_all_endpoints_when_later_endpoint_rejects_flag( "receiver_transfer_engine_infos": [{"tp_rank": 0, "session_id": "recv-0:7000"}], }, ) - ep0_full_response = _FakeResponse( - 200, - { - "success": True, - "message": "full", - "tensor_map": {param_name: [ep0_full_locator]}, - "receiver_transfer_engine_infos": [{"tp_rank": 0, "session_id": "recv-0:7000"}], - }, - ) ep1_full_response = _FakeResponse( 200, { @@ -547,26 +600,24 @@ def test_cached_prepare_restarts_all_endpoints_when_later_endpoint_rejects_flag( with patch( "requests.post", - side_effect=[ep0_cached_response, _FakeResponse(422, {}), ep0_full_response, ep1_full_response], + side_effect=[ep0_cached_response, _FakeResponse(422, {}), ep1_full_response], ) as posted: ok = backend.initialize() assert ok is True - assert posted.call_count == 4 + 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-0:5000/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"] - assert "p2p_return_tensor_map" not in posted.call_args_list[3].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": 0x1000, + "recv-0:7000": 0xDEAD0000, "recv-1:7000": 0x2000, } assert backend._receiver_session_ids == ["recv-0:7000", "recv-1:7000"] @@ -600,6 +651,16 @@ def test_complete_sync_preserves_cached_prepare_state(self): 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): @@ -768,6 +829,47 @@ def test_transfer_bucket_writes_correct_slice_per_receiver(self): 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 = {} @@ -1063,6 +1165,7 @@ def test_transfer_bucket_does_not_coalesce_without_receiver_memory_handles(self, 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 = { @@ -1101,11 +1204,102 @@ def test_transfer_bucket_failure_names_tensor_and_receiver_handle(self, monkeypa 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 + 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", @@ -1114,15 +1308,7 @@ def _make_multi_sender_backend( buffer_size_mb=64, device="cuda:0", training_rank=rank_index, - 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, - }, + backend_config=backend_config, ) backend = P2PTransportBackend(cfg) fake_engine = _FakeMooncakeEngine(session_id=f"trainer-{rank_index}:1234") @@ -1135,6 +1321,373 @@ def test_supports_direct_ep_transfer_advertises_all_ranks(self): 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) @@ -1166,7 +1719,7 @@ def initialize_rank0(): all_gather_calls = [] - def all_gather_object(output, value): + def all_gather_object(output, value, **kwargs): all_gather_calls.append(value) if len(all_gather_calls) == 1: output[:] = [False, False] @@ -1186,6 +1739,105 @@ def all_gather_object(output, value): 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) @@ -1337,6 +1989,28 @@ def test_complete_sync_defaults_flush_cache_false(self): 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): From e8fcbf09e0265c2334d3a0b18d8f5a28ac631338 Mon Sep 17 00:00:00 2001 From: Ashwinee Panda Date: Wed, 20 May 2026 16:02:50 -0700 Subject: [PATCH 38/49] feat: Add on-policy distillation MVP MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Add on-policy distillation MVP * Add SGLang teacher OPD e2e test * Add OPD full-rollout e2e test with student/teacher SGLang servers * Fix nccl_broadcast TCPStore master deadlock Master TCPStore was created with the PyTorch default wait_for_workers=True, which blocks the constructor until world_size-1 workers connect. The design in init_nccl_group is to create a non-blocking listening master first, then fire /init_weights_update_group at the inference endpoints (workers) from a background thread, then complete the NCCL rendezvous in the main thread via _new_process_group_helper. With the default, the master constructor blocks before the background thread can even start, so workers never receive the HTTP request and sync_inference_weights deadlocks until the request times out. Pass wait_for_workers=False so the master returns as soon as it is listening; the actual rendezvous still synchronizes via _init_training_process_group. Update the unit-test fake TCPStore to accept extra kwargs. * Document NCCL sync hang in OPD full-pipeline test docstring The test now keeps the disk-based weight refresh path (xorl save_weights_for_sampler + SGLang update_weights_from_disk) as the primary validation. Adding a NOTE explaining that the NCCL path was also wired up in this branch: - The wait_for_workers=False fix (separate commit) unblocks the sync_inference_weights call past TCPStore master construction. - With the fix in place, the broadcast itself now succeeds β€” the trainer reaches "Transfer complete: 0.02s, 5 buckets, 26 params". - The cleanup phase still hangs: SGLang's /destroy_weights_update_group never returns for the tiny-model setup, which causes dist.destroy_process_group on the trainer side to hang as well. That's an SGLang side issue and is out of scope here. * Add multi-pod OPD k8s pipeline (student SGLang / teacher SGLang / trainer) Three-pod cross-node validation of the full OPD loop, so the NCCL weight sync exercises the actual cluster fabric instead of co-locating everything on the same node: experiments/local_benchmark/k8s/opd_full_pipeline/ student-sglang.yaml # SGLang sampler, registered as inference endpoint teacher-sglang.yaml # SGLang with --enable-return-hidden-states trainer.yaml # xorl.server.launcher + run_opd_pipeline.py server_config.yaml # 1-GPU bf16 Qwen3-0.6B server config scripts/opd/ run_opd_pipeline.py # rollout -> teacher hidden -> fwd_bwd -> optim -> sync submit_k8s_pipeline.sh Pods publish their endpoints to /shared/opd-coord//{student,teacher}.json. The trainer driver reads those, registers the student via /add_inference_endpoint, and runs OPD_NUM_STEPS iterations of forward_backward(opd_loss) + optim_step + sync_inference_weights, asserting each broadcast succeeds and rollouts diverge. podAntiAffinity guarantees the three roles land on distinct nodes; nodeAffinity excludes hosts known to have broken IB / HBM ECC. * Driver: verify post-sync rollout instead of trusting sync_inference_weights HTTP status SGLang's /destroy_weights_update_group does not respond after the NCCL weight broadcast completes (the trainer log reports Transfer complete in ~0.2 s, but the destroy POST stalls at the TCP read), so the trainer's sync_inference_weights call eventually returns 500. The broadcast itself has already updated the inference weights, so verify success by sampling greedy rollouts before and after the sync and checking they diverged. Also bump the bucket size to 4 GB so all params land in one bucket and fewer HTTP/NCCL hops are needed per sync. * k8s OPD: lower mem fraction, raise backoffLimit for tenant admission races * k8s OPD: drop done marker, keep SGLang pods alive, fix trainer-retry deadlock * k8s OPD: trainer retries up to 8x and fast-fails on OOM-during-launch * k8s OPD: longer trainer health probe and tolerate either JSON spacing * k8s OPD: trainer fast-fails when worker OOMs but launcher survives * weight-sync: abort training pg on destroy + disable teacher radix cache Two fixes pulled together: * nccl_broadcast.destroy_nccl_group: replace dist.destroy_process_group(self.process_group) with self.process_group.abort(). destroy_process_group calls pg.shutdown(), which cooperatively waits for the inference peer to finalize. With the matching SGLang fix (), the inference side now aborts its comm immediately, so a cooperative shutdown on the trainer side hangs until the 600 s engine timeout fires. Aborting on both sides is the right semantic when the broadcast has already completed β€” the comm is done, we just want to tear it down. * opd k8s teacher SGLang: pass --disable-radix-cache. With radix caching on, /generate's prefill hidden states are returned only for tokens not already cached, so step 2 of the OPD loop sees a [seq_len-prefix]-shaped tensor for an input of seq_len tokens and the driver fails. Disabling the cache forces a fresh prefill per request. * Driver: assert post-sync rollout diverges from *initial* (not previous step) After the SGLang+xorl abort fixes the pipeline runs end-to-end, but with the tiny Qwen3-0.6B model two OPD steps drive lm_head into a degenerate state that emits the same constant rollout every iteration. The previous assertion compared post-sync rollout to *the previous step's* post-sync rollout, which fails when both are degenerate even though the very first sync demonstrably changed the weights. Snapshot the initial greedy rollout once and assert that at least one step's post-sync rollout differs from it. * lint: ruff format on touched files + whitelist 'NotIn' in codespell Brings the OPD branch through the repo's pre-commit hooks: ruff-format reflows a few long lines and codespell stops flagging NotIn (the k8s nodeAffinity operator). * k8s OPD: replace personal $HOME paths with /workspace placeholders tests/test_example_assets.py forbids /home// in tracked example and experiment assets. Switch the OPD manifests to /workspace/home and /workspace/{xorl,xorl-sglang}, matching the convention in the existing qwen3-32b-muon-1node-job.yaml and friends. The home-apanda PVC name is the cluster resource and stays. * fix: OPD loss reporting β€” fp32 across SP ranks to keep all_reduce sane With ulysses sequence sharding, ranks holding only IGNORE_INDEX tokens hit the OPD early-return paths and produced a bf16 zero-loss while the rank actually computing the KL returned fp32 from total_weighted_kl / denom. NCCL's dist.all_reduce silently reinterprets the bytes when participating tensors disagree on dtype, turning the per-step reported loss into garbage on the rank-0 reporter (observed: -1.99e+35 with finite gradients on the compute rank). Three converging fixes: - _zero_loss_with_graph casts inputs to fp32 before * 0.0. - _compute_opd_micro_batch_loss uses fp32 in both the early-return branch and the total_loss accumulator initializer. - forward_backward's loss reporter casts loss_report to fp32 before the cross-rank reduction as a defensive measure. Also makes max_adapters_per_model environment-overridable (XORL_MAX_ADAPTERS_PER_MODEL, default 32) so multi-step OPD runs don't have to evict and trigger SGLang's flaky /unload_lora_adapter. Validated end-to-end on the Countdown 24-puzzle 1-pod manifest: 16/16 OPD steps clean, loss 1.29 -> 0.58, exact 3/32 -> 8/32. * review: PP guard, fail-loud cache indices, drop misleading non_blocking Address review comments on: - Add a symmetric `if self.pp_enabled: raise NotImplementedError(...)` inside _compute_opd_micro_batch_loss matching the existing TP guard. The dispatcher-level guard in _run_forward_backward / _run_forward already blocks PP, but this prevents a future direct caller from silently launching OPD under a parallel topology that hasn't been thought through. - Replace `flat_indices.clamp(min=0)` in TeacherActivationCache.get with an explicit IndexError on negative values. The clamp masked producer bugs (an off-by-one in teacher_cache_indices construction was caught only via the loss magnitude going wild). Add tests for both negative and out-of-range index paths. - Drop `non_blocking=True` on the.to() calls in TeacherHeadManager.get and TeacherActivationCache.get β€” the source CPU tensors come from safetensors / torch.load without `pin_memory()`, so the flag is a no-op and misleads future readers about the copy behavior. - Document why TeacherHeadManager keeps only one teacher head on device at a time (vocab * hidden bytes per teacher; multi-teacher workloads can revisit with an LRU later). * Address OPD review comments * Fix OPD teacher hidden state handling * feat(weight-sync): add Mooncake P2P backend * test(weight-sync): add sync stress harness * fix(weight-sync): fail closed on P2P transfer errors * fix: correct OPD metrics and shifted teacher fields * fix(weight-sync): pin mooncake and harden p2p routing * fix(weight-sync): return measured transfer time * test(weight-sync): report sync timing breakdown * fix(models): expand qwen3.5 layer types * fix(weight-sync): prefer fast endpoint health probe * test: relax OLMo2 checkpoint tolerance * test: relax OLMo2 logits tolerance * fix: expand Qwen3.5 MoE layer types * fix(weight-sync): stabilize async p2p medium batches * fix(weight-sync): fail closed on p2p finalization * fix(weight-sync): clean up p2p diagnostics * fix(weight-sync): speed up p2p fp8 formatting * fix(weight-sync): stream fp8 cpu workspace transfers * feat: add OPD teacher profiling path * style: format OPD profiling changes * chore: mark OPD teacher-store helper executable * chore: make OPD k8s manifests path-generic * docs: record OPD profile trial results * docs: fix OPD trial log formatting * fix(weight-sync): honor remote sync timeout * test: run OPD weight sync profiling path * fix: launch multinode workers through python module * fix: launch workers without torchrun wrapper * fix(weight-sync): bound fp8 cpu workspace staging * test: record OPD weight sync trials * chore(weight-sync): add sparse delta probe * fix(weight-sync): handle qwen3.5 p2p linear attention slices * test: record OPD p2p retry after qwen35 sync fix * chore(weight-sync): mark delta probe executable * fix(weight-sync): require receiver handles for p2p coalescing * test: record OPD p2p receiver registration retry * test: record OPD p2p locator registration retry * test: record OPD p2p allocator registration retry * test: record OPD p2p location registration retry * test: record OPD p2p segment registration retry * test: record OPD no-expand p2p retry * fix(weight-sync): support kimi p2p receiver layouts * Fix P2P weight sync edge cases * fix(weight-sync): restart cached p2p prepare on fallback * fix(weight-sync): address p2p fp8 review issues * fix(weight-sync): handle qwen3.5 p2p linear attention slices * Record OPD P2P CPU staging retry * fix server adapter and metric regressions * fix(weight-sync): complete qwen3.5 p2p opd sync * Revert "fix server adapter and metric regressions" This reverts commit fabd8ae3b3b7ccb9af3ac6a079b7e10fc067ea9d. * fix(server): update adapter, routing, and p2p handling * fix(opd): correct distributed teacher hidden cache writes * fix(opd): stabilize fsdp optimizer smoke Add the OPD profiling pipeline updates used for the async smoke runs, anchor lm_head.forward for FSDP OPD loss, and synchronize after optimizer.step() before releasing gradient storage. Validation: PYTHONPATH=src PYTHONDONTWRITEBYTECODE=1 pytest -q tests/server/runner/test_opd_runner.py tests/ops/loss/test_opd_loss.py tests/scripts/test_opd_pipeline_payloads.py * fix(opd): require post-sync version acknowledgement * fix(opd): support p2p sync for tied lm head * fix(opd): scope NCCL sync groups per request * fix(merge): align optim_step LR resolution, IS metrics, and adapter optimizer plumbing CPU-test fixes after the mainβ†’ merge: - training_ops: keep HEAD's _optim_step_learning_rate helper (registered sessions fall back to AdamParams default) but extend the chain with main's _server_default_learning_rate and raise loudly when the session is not registered, satisfying both the HEAD fallback test and main's fail-loud test. - model_runner: restore main's full _accumulate_is_metrics / _finalize_is_metrics that store Python scalars (via _metric_to_float) so the IS-only metric tests pass; HEAD's _accumulate_loss_metrics still drives the OPD path. - adapters/manager: thread cautious_weight_decay (and the rest of the manager-level optimizer_config) through the legacy session spec so optimizer_config={'cautious_weight_decay': True} reaches build_optimizer. - session_spec: propagate cautious_weight_decay in session_optimizer_build_kwargs. - request_processor: only inject routed_experts / routed_expert_logits for forward_backward so non-MoE backends (e.g. TeacherCacheCPUBackend) do not get unexpected kwargs. * fix(merge): always thread routed_experts/logits through model_pass kwargs The orchestrator request_processor now passes routed_experts and routed_expert_logits to both forward and forward_backward backend methods, matching the multi-LoRA replay contract introduced in. The OPD test backend (TeacherCacheCPUBackend.forward) is updated to accept the new kwargs so its forward signature is compatible with the unconditional pass-through. --------- Co-authored-by: Qingyang Wu --- .../data/collators/packing_concat_collator.py | 41 +- .../data/collators/sequence_shard_collator.py | 50 +- src/xorl/data/collators/tensor_collator.py | 5 +- src/xorl/distillation/__init__.py | 11 + src/xorl/distillation/teacher_cache.py | 299 ++++++ src/xorl/distillation/teacher_store.py | 366 +++++++ src/xorl/models/module_utils.py | 2 + src/xorl/ops/loss/__init__.py | 4 + src/xorl/ops/loss/compiled_cross_entropy.py | 87 ++ src/xorl/ops/loss/opd_loss.py | 214 +++++ src/xorl/ops/loss/opd_streaming_kl.py | 188 ++++ src/xorl/server/api_server/server.py | 15 +- src/xorl/server/api_server/training_ops.py | 75 +- src/xorl/server/launcher.py | 29 +- src/xorl/server/orchestrator/packing.py | 74 +- .../server/orchestrator/request_processor.py | 66 +- src/xorl/server/runner/adapters/manager.py | 26 +- src/xorl/server/runner/model_runner.py | 909 +++++++++++++++++- src/xorl/server/runner/runner_dispatcher.py | 35 +- src/xorl/server/runner/utils/batch_utils.py | 85 +- src/xorl/server/session_spec.py | 3 + .../weight_sync/backends/nccl_broadcast.py | 14 +- src/xorl/server/weight_sync/backends/p2p.py | 23 +- tests/_helpers/__init__.py | 0 tests/_helpers/opd.py | 126 +++ .../collators/test_packing_concat_collator.py | 31 + .../collators/test_sequence_shard_collator.py | 32 + tests/e2e/e2e_utils.py | 14 +- tests/e2e/server_utils.py | 50 +- tests/e2e/test_opd_cpu.py | 246 +++++ tests/e2e/test_opd_full_pipeline.py | 450 +++++++++ tests/e2e/test_opd_gpu.py | 78 ++ tests/e2e/test_opd_sglang_teacher.py | 245 +++++ tests/models/test_module_utils_broadcast.py | 8 + tests/models/test_olmo2_support.py | 4 +- tests/ops/loss/test_opd_loss.py | 231 +++++ tests/server/api_server/test_api_server.py | 45 + tests/server/orchestrator/test_packing.py | 74 ++ .../orchestrator/test_request_processor.py | 40 + tests/server/runner/test_batch_utils.py | 68 ++ .../runner/test_lora_checkpoint_roundtrip.py | 19 + tests/server/runner/test_opd_runner.py | 403 ++++++++ tests/server/runner/test_runner_dispatcher.py | 106 ++ .../weight_sync/test_p2p_backend_protocol.py | 50 + .../utils/test_distillation_teacher_cache.py | 137 +++ 45 files changed, 4929 insertions(+), 149 deletions(-) create mode 100644 src/xorl/distillation/__init__.py create mode 100644 src/xorl/distillation/teacher_cache.py create mode 100644 src/xorl/distillation/teacher_store.py create mode 100644 src/xorl/ops/loss/opd_loss.py create mode 100644 src/xorl/ops/loss/opd_streaming_kl.py create mode 100644 tests/_helpers/__init__.py create mode 100644 tests/_helpers/opd.py create mode 100644 tests/e2e/test_opd_cpu.py create mode 100644 tests/e2e/test_opd_full_pipeline.py create mode 100644 tests/e2e/test_opd_gpu.py create mode 100644 tests/e2e/test_opd_sglang_teacher.py create mode 100644 tests/ops/loss/test_opd_loss.py create mode 100644 tests/server/runner/test_batch_utils.py create mode 100644 tests/server/runner/test_opd_runner.py create mode 100644 tests/server/runner/test_runner_dispatcher.py create mode 100644 tests/utils/test_distillation_teacher_cache.py 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/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/models/module_utils.py b/src/xorl/models/module_utils.py index dd9e3096..bee30f69 100644 --- a/src/xorl/models/module_utils.py +++ b/src/xorl/models/module_utils.py @@ -341,6 +341,8 @@ 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 diff --git a/src/xorl/ops/loss/__init__.py b/src/xorl/ops/loss/__init__.py index 501064b8..cd9bd894 100644 --- a/src/xorl/ops/loss/__init__.py +++ b/src/xorl/ops/loss/__init__.py @@ -14,6 +14,7 @@ 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 @@ -33,6 +34,7 @@ "importance_sampling": importance_sampling_loss_function, "policy_loss": policy_loss_function, "drgrpo": drgrpo_loss_function, + "opd_loss": opd_loss_function, } @@ -51,6 +53,7 @@ def register_loss_function(name: str, fn: Callable) -> None: __all__ = [ "CrossEntropyMode", "LossOutput", + "OPDLossMetrics", "LOSS_REGISTRY", "Reducer", "SequencePartial", @@ -60,6 +63,7 @@ def register_loss_function(name: str, fn: Callable) -> None: "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/compiled_cross_entropy.py b/src/xorl/ops/loss/compiled_cross_entropy.py index 2e74868e..a62d8c58 100644 --- a/src/xorl/ops/loss/compiled_cross_entropy.py +++ b/src/xorl/ops/loss/compiled_cross_entropy.py @@ -17,6 +17,9 @@ # 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 @@ -102,6 +105,76 @@ def compiled_ce_and_lse_sq_function( 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. @@ -158,3 +231,17 @@ def _compute_ce_and_lse_sq(hidden_states, weight, labels, ignore_index): 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/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/server/api_server/server.py b/src/xorl/server/api_server/server.py index 47fc30db..a973ea82 100644 --- a/src/xorl/server/api_server/server.py +++ b/src/xorl/server/api_server/server.py @@ -171,8 +171,11 @@ def __init__( # 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()) @@ -308,9 +311,13 @@ 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, *, notify_workers: bool = True) -> None: """ diff --git a/src/xorl/server/api_server/training_ops.py b/src/xorl/server/api_server/training_ops.py index 85a9aea5..b2048cff 100644 --- a/src/xorl/server/api_server/training_ops.py +++ b/src/xorl/server/api_server/training_ops.py @@ -10,6 +10,7 @@ from fastapi import HTTPException, status from xorl.server.api_server.api_types import ( + AdamParams, ForwardBackwardRequest, ForwardBackwardResponse, ForwardRequest, @@ -70,6 +71,41 @@ def _server_default_learning_rate(self) -> Optional[float]: 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 # ========================================================================= @@ -264,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) @@ -282,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 @@ -375,6 +411,15 @@ 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, @@ -407,21 +452,7 @@ async def optim_step(self, request: OptimStepRequest) -> OptimStepResponse: try: adam_params = request.adam_params - lr = request.learning_rate - if lr is None and adam_params is not None: - lr = adam_params.learning_rate - if lr is None: - lr = self._session_default_learning_rate(request.model_id) - if lr is None: - lr = self._server_default_learning_rate() - if lr is None: - 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" - ), - ) + lr = self._optim_step_learning_rate(request) # Determine gradient clipping value # Priority: explicit gradient_clip parameter, then adam_params.grad_clip_norm diff --git a/src/xorl/server/launcher.py b/src/xorl/server/launcher.py index 010788bb..a7dbed69 100644 --- a/src/xorl/server/launcher.py +++ b/src/xorl/server/launcher.py @@ -127,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}") # ============================================================================ @@ -914,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}", diff --git a/src/xorl/server/orchestrator/packing.py b/src/xorl/server/orchestrator/packing.py index 567245f7..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.) @@ -812,11 +869,22 @@ def _add_sample_to_batch( 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 353ee8fa..e89f6a0d 100644 --- a/src/xorl/server/orchestrator/request_processor.py +++ b/src/xorl/server/orchestrator/request_processor.py @@ -164,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( @@ -196,8 +206,8 @@ async def _execute_model_pass( loss_fn=loss_fn, loss_fn_params=loss_fn_params, model_id=p.model_id, - routed_experts=p.routed_experts, - routed_expert_logits=p.routed_expert_logits, + routed_experts=routed_experts, + routed_expert_logits=routed_expert_logits, request_id=request.request_id, ) @@ -213,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 @@ -230,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) @@ -268,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. @@ -572,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 [ { @@ -595,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, diff --git a/src/xorl/server/runner/adapters/manager.py b/src/xorl/server/runner/adapters/manager.py index c0a339da..d4ca5afc 100644 --- a/src/xorl/server/runner/adapters/manager.py +++ b/src/xorl/server/runner/adapters/manager.py @@ -353,22 +353,30 @@ def _legacy_session_spec(self, *, lr: float) -> Dict[str, Any]: 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) - 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": { + # 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(self.weight_decay), + "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: diff --git a/src/xorl/server/runner/model_runner.py b/src/xorl/server/runner/model_runner.py index b03385fa..ebb9828b 100644 --- a/src/xorl/server/runner/model_runner.py +++ b/src/xorl/server/runner/model_runner.py @@ -25,10 +25,12 @@ import torch import torch.distributed as dist import torch.nn.functional as F +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 @@ -37,9 +39,12 @@ 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 @@ -192,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__( @@ -249,6 +280,12 @@ def __init__( # 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]] = {} @@ -937,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: @@ -961,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: @@ -984,6 +1021,481 @@ def _collect_per_token_outputs(self, per_token_tensors, micro_batch, accumulator if gathered.get("loss") is not None: accumulators["losses"].append(gathered["loss"].cpu()) + @staticmethod + 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 = 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.""" @@ -1077,10 +1589,271 @@ def _count_active_microbatches(self, micro_batches) -> tuple[int, int]: 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.""" params = loss_fn_params or {} @@ -1092,7 +1865,10 @@ def _compute_micro_batch_loss(self, micro_batch, loss_fn, loss_fn_params): outputs = self.model(**model_inputs, use_cache=False, output_hidden_states=False) hidden_states = outputs.last_hidden_state effective_weight = self._get_effective_lm_head_weight() - metric_sum_reducer = TokenPartial(scale=hidden_states.new_tensor(1.0, dtype=torch.float32)) + + # 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 @@ -1108,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 @@ -1128,9 +1904,9 @@ 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=metric_sum_reducer, + 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 @@ -1234,9 +2010,9 @@ 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=metric_sum_reducer, + 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 @@ -1248,10 +2024,22 @@ def _compute_micro_batch_loss(self, micro_batch, loss_fn, loss_fn_params): 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, is_metric_ops, outputs + return local_loss_sum, per_token_outputs, is_metrics, is_metric_ops, outputs # ========================================================================= # Per-sample K3 KL divergence @@ -1322,6 +2110,15 @@ def _forward_loop( 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: @@ -1333,8 +2130,20 @@ def _forward_loop( 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) @@ -1359,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, is_metric_ops, 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 @@ -1424,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") @@ -1433,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, is_metric_ops) + # 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, 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). @@ -1482,11 +2298,27 @@ def _forward_loop( 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"]] @@ -1503,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 @@ -1638,6 +2470,9 @@ def forward_backward( 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 compute_ref_logprobs = params.get("compute_ref_logprobs", False) @@ -1790,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) @@ -1854,6 +2706,9 @@ def forward( 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( @@ -1990,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: diff --git a/src/xorl/server/runner/runner_dispatcher.py b/src/xorl/server/runner/runner_dispatcher.py index 0d62b2ce..c183007d 100644 --- a/src/xorl/server/runner/runner_dispatcher.py +++ b/src/xorl/server/runner/runner_dispatcher.py @@ -724,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. @@ -748,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) 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/session_spec.py b/src/xorl/server/session_spec.py index 0caf2d1c..bc52067d 100644 --- a/src/xorl/server/session_spec.py +++ b/src/xorl/server/session_spec.py @@ -311,6 +311,9 @@ def session_optimizer_build_kwargs(optimizer_config: Dict[str, Any]) -> Dict[str 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 diff --git a/src/xorl/server/weight_sync/backends/nccl_broadcast.py b/src/xorl/server/weight_sync/backends/nccl_broadcast.py index 63ad675b..c06929e9 100644 --- a/src/xorl/server/weight_sync/backends/nccl_broadcast.py +++ b/src/xorl/server/weight_sync/backends/nccl_broadcast.py @@ -366,9 +366,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() diff --git a/src/xorl/server/weight_sync/backends/p2p.py b/src/xorl/server/weight_sync/backends/p2p.py index ddb171a0..43a32186 100644 --- a/src/xorl/server/weight_sync/backends/p2p.py +++ b/src/xorl/server/weight_sync/backends/p2p.py @@ -1724,6 +1724,9 @@ def transfer_bucket( 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 @@ -2322,6 +2325,18 @@ def _locators_for_source_name(self, name: str) -> Optional[List[Dict[str, Any]]] # 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, @@ -2519,9 +2534,11 @@ def _resolve_local_hostname() -> str: Mooncake's handshake binds on this hostname; it must be reachable from the SGLang receiver. """ - explicit = os.environ.get("XORL_P2P_HOSTNAME") - if explicit and explicit.strip(): - return explicit.strip() + # 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: 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/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/server_utils.py b/tests/e2e/server_utils.py index f8177fd7..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 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_module_utils_broadcast.py b/tests/models/test_module_utils_broadcast.py index 584d4f2a..e9cd2815 100644 --- a/tests/models/test_module_utils_broadcast.py +++ b/tests/models/test_module_utils_broadcast.py @@ -493,6 +493,14 @@ def fake_try_load_state_dict_local(weights_path, **kwargs): 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 = [] diff --git a/tests/models/test_olmo2_support.py b/tests/models/test_olmo2_support.py index 5ea04789..9791dd6f 100644 --- a/tests/models/test_olmo2_support.py +++ b/tests/models/test_olmo2_support.py @@ -176,5 +176,5 @@ def test_olmo2_checkpoint_handler_loads_hf_weights_into_fused_model(): hf_logits = hf_model.lm_head(hf_hidden_states) xorl_logits = xorl_model.lm_head(xorl_hidden_states) - torch.testing.assert_close(xorl_hidden_states, hf_hidden_states, atol=5e-5, rtol=5e-5) - torch.testing.assert_close(xorl_logits, hf_logits, atol=5e-5, rtol=5e-5) + 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/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/server/api_server/test_api_server.py b/tests/server/api_server/test_api_server.py index d7be18ae..719567f4 100644 --- a/tests/server/api_server/test_api_server.py +++ b/tests/server/api_server/test_api_server.py @@ -170,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.""" diff --git a/tests/server/orchestrator/test_packing.py b/tests/server/orchestrator/test_packing.py index fb93e40f..beb00447 100644 --- a/tests/server/orchestrator/test_packing.py +++ b/tests/server/orchestrator/test_packing.py @@ -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) diff --git a/tests/server/orchestrator/test_request_processor.py b/tests/server/orchestrator/test_request_processor.py index ffdc8de2..fc36d555 100644 --- a/tests/server/orchestrator/test_request_processor.py +++ b/tests/server/orchestrator/test_request_processor.py @@ -33,6 +33,7 @@ OptimStepData, RegisterSessionData, SaveStateData, + SyncWeightsData, ) from xorl.server.runner.runner_dispatcher import RunnerDispatcher @@ -112,6 +113,45 @@ 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.""" 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_lora_checkpoint_roundtrip.py b/tests/server/runner/test_lora_checkpoint_roundtrip.py index e331c397..ec1b45fe 100644 --- a/tests/server/runner/test_lora_checkpoint_roundtrip.py +++ b/tests/server/runner/test_lora_checkpoint_roundtrip.py @@ -16,6 +16,7 @@ save_lora_checkpoint, ) from xorl.models.layers.moe import MoEExpertsLoRA, MoELoRAConfig +from xorl.optim import AnyPrecisionAdamW pytestmark = [pytest.mark.cpu, pytest.mark.server] @@ -348,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_opd_runner.py b/tests/server/runner/test_opd_runner.py new file mode 100644 index 00000000..ea2d4c87 --- /dev/null +++ b/tests/server/runner/test_opd_runner.py @@ -0,0 +1,403 @@ +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) + + +@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.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_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/weight_sync/test_p2p_backend_protocol.py b/tests/server/weight_sync/test_p2p_backend_protocol.py index 314feeb1..40ff2798 100644 --- a/tests/server/weight_sync/test_p2p_backend_protocol.py +++ b/tests/server/weight_sync/test_p2p_backend_protocol.py @@ -199,6 +199,19 @@ def _hf_locator( 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) @@ -878,6 +891,43 @@ def test_transfer_bucket_fails_on_unknown_param(self): 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" 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") From 48c33d41b641eb79c526102aef004b5aca44815e Mon Sep 17 00:00:00 2001 From: Ashwinee Panda Date: Thu, 21 May 2026 10:45:36 -0700 Subject: [PATCH 39/49] fix(server): R3 routing replay picks Qwen3.6 nested top-k + model_topk first * fix(server): R3 routing replay picks Qwen3.6 nested top-k + model_topk first `_extract_topk` now checks `config.text_config.num_experts_per_tok` before `config.num_experts_per_tok` so Qwen3.6 (which nests the field under text_config) does not silently miss it. In `fill_routing_replay`, prefer `self._model_topk` over inferring from `decoded_routing[0][0][0]`. With mixed-width rows (e.g. Qwen3.6 returns some 6-wide and some 8-wide tokens), the previous code would pick row 0's width and crash tensorization on later rows. * chore(lint): apply ruff-format to routing_replay_handler.py --- .../runner/utils/routing_replay_handler.py | 29 ++++++++++++++++--- 1 file changed, 25 insertions(+), 4 deletions(-) diff --git a/src/xorl/server/runner/utils/routing_replay_handler.py b/src/xorl/server/runner/utils/routing_replay_handler.py index 1e13c12b..c107d45e 100644 --- a/src/xorl/server/runner/utils/routing_replay_handler.py +++ b/src/xorl/server/runner/utils/routing_replay_handler.py @@ -67,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 @@ -254,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( From 690cfa5d0b3ae9ba36a3e30df1aed08938e51f01 Mon Sep 17 00:00:00 2001 From: Ashwinee Panda Date: Thu, 21 May 2026 12:50:26 -0700 Subject: [PATCH 40/49] feat(skill): add XORL throughput tuner * Add XORL throughput tuner skill * Update XORL throughput tuner repro guidance * Fix renderer provenance block indentation * Add Qwen3.6 benchmark recipe * Add throughput tuner skill tests * feat: add K3 correctness gating artifacts * fix: harden K3 replay scheduling * docs: record Qwen3.5 K3 replay attempt * docs: record Qwen3.5 K3 gate pass * feat(moe): add XORL_MOE_SYNTHETIC_ROUTING=balanced diagnostic knob Adds an env-driven synthetic routing mode that overrides TopKRouter.forward and the K3 replay path with a deterministic balanced top-k pattern (experts cycling [0..top_k-1], [top_k..2*top_k-1],...). Lets the throughput tuner isolate routing-induced load imbalance from real-model training cost. K3 static-trace replay correctly unsets the env so real-traffic traces are not overwritten. * feat(bench): add Q3.5-397B R73/R75 shortctx sweep configs Colocates the 12 qwen3_5_397b_a17b shortctx-8node configs next to the existing R73 K3 gate + shortctx MFU summary under skills/xorl-throughput-tuner/benchmarks/qwen3_5_397b_a17b/configs/. Covers R73 (no-async compile + activation-offload prefetch=4) and R75 (async compile) winning recipes plus R60 ns3, R61 mbs2 offload, R63 alltoall, R69 mbs4, R74 mbs5 prefetch=2 ablations. Also fixes copy-paste bugs from the original sweep generator: r69 wandb_name r67 -> r69, tag nocompile -> compile r73 wandb_name r70/prefetch2 -> r73/prefetch4, tag prefetch2 -> prefetch4 r74 wandb_name r70 -> r74 + _async suffix, tag noasync -> async Supersedes. * fix(skill): address open review comments on throughput tuner Bug: Add qwen3.6-35b-a3b to MODEL_PRESETS so the bundled benchmark recipe hits the EP-divides-experts and ulysses-vs-KV-heads guardrails instead of silently falling through to None. Nits: Producer (run_local_benchmark.py) STEP_RE now captures the peak_mem field the trainer already emits, so benchmark_summary.json carries peak_mem_gb per step; parse_summary in collect_xorl_metrics no longer reports peak_mem_gb_max=None on real runs. discover() preserves log-only sibling runs when other runs in the same sweep tree have benchmark_summary.json (was all-or-nothing per input path). Default use_wandb to false in the portable qwen3_6_35b_a3b benchmark config; the renderer already passes WANDB_MODE/WANDB_ENTITY through for users on clusters that have credentials. Replace brittle 'value: "2"' substring test with a regex check that asserts NCCL_NET_GDR_LEVEL stays as PHB. Make securityContext.privileged a --privileged/--no-privileged flag (default true) so stricter clusters can render without it. Test coverage: Add tests for parse_summary peak_mem plumbing, discover() per-dir preference, the qwen3.6 preset resolution path, and the --no-privileged renderer mode. --- src/xorl/models/layers/moe/moe_block.py | 12 ++++- src/xorl/models/layers/moe/router.py | 41 +++++++++++++++++ .../models/test_moe_train_router_dispatch.py | 46 +++++++++++++++++++ 3 files changed, 98 insertions(+), 1 deletion(-) diff --git a/src/xorl/models/layers/moe/moe_block.py b/src/xorl/models/layers/moe/moe_block.py index 1bd36e0f..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 @@ -123,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: 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/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)) From dd5553584013b7f17179ad0d4edeb2a0977718a0 Mon Sep 17 00:00:00 2001 From: Ashwinee Panda Date: Thu, 21 May 2026 18:06:23 -0700 Subject: [PATCH 41/49] fix(deepep): default to synchronous combine, env-gate unsafe async * fix(deepep): default to synchronous combine, env-gate unsafe async The MoE output from tokens_post_combine is consumed by the transformer block before the next DeepEP dispatch, so deferring the wait lets downstream compute read incomplete combine data. Force async_combine off in tokens_post_combine unless XORL_DEEPEP_UNSAFE_ASYNC_COMBINE is explicitly set. Keep the API flag for provenance. * Address review: warn on neutered async_combine, cache env, update docs - Cache XORL_DEEPEP_UNSAFE_ASYNC_COMBINE at module import (_ALLOW_UNSAFE_ASYNC_COMBINE) instead of re-reading os.environ per layer call, matching the _DEEPEP_PROFILE pattern. - Log a one-time warning when async_combine=True is silently demoted, so users who set deepep_async_combine in YAML are not left wondering why no overlap occurs. - Update tokens_post_combine docstring to describe the env gate and the underlying default-stream race instead of the original async-path contract. - Update config-reference, expert_parallelism, and moe/deepep docs to flag that the async path is disabled by default and how to opt in. - Mark the guard test as cpu and monkeypatch the cached constant rather than the env. * Address review: simplify env check to == "1" Match the codebase convention used in handler.py for env-var boolean flags. --- .../content/docs/config-reference/local.md | 2 +- .../content/docs/config-reference/server.md | 2 +- docs/src/content/docs/moe/deepep.mdx | 24 ++++++-- .../docs/parallelism/expert_parallelism.mdx | 19 +++++-- src/xorl/distributed/moe/deepep.py | 27 +++++++-- .../test_deepep_async_combine_guard.py | 55 +++++++++++++++++++ 6 files changed, 113 insertions(+), 16 deletions(-) create mode 100644 tests/distributed/test_deepep_async_combine_guard.py diff --git a/docs/src/content/docs/config-reference/local.md b/docs/src/content/docs/config-reference/local.md index 12ba3a83..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). | diff --git a/docs/src/content/docs/config-reference/server.md b/docs/src/content/docs/config-reference/server.md index e376c166..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). | 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/expert_parallelism.mdx b/docs/src/content/docs/parallelism/expert_parallelism.mdx index 80cc4903..eabc55af 100644 --- a/docs/src/content/docs/parallelism/expert_parallelism.mdx +++ b/docs/src/content/docs/parallelism/expert_parallelism.mdx @@ -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 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/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 From b4150e5f8c8679badfa12aafc9d9cd79813deaba Mon Sep 17 00:00:00 2001 From: Ashwinee Panda Date: Fri, 22 May 2026 12:35:07 -0700 Subject: [PATCH 42/49] fix(qwen3.5): mrope_interleaved must not pairwise-rotate q/k features MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(qwen3.5): mrope_interleaved must not pairwise-rotate q/k features The interleaved flag controls how T/H/W frequency sections are mixed when MRoPE builds cos/sin, not how q/k features are paired. Match HF/SGLang's standard half-rotate. Previously we rebuilt cos/sin to an interleaved layout and called a pairwise rotate_half, which diverged from the reference implementation on Qwen3.5/Qwen3.6. * Narrow fix to Qwen3.5/3.6 modeling β€” preserve DSv3 pairwise path The previous revision made `qwen3_5_apply_rotary_pos_emb` ignore its `interleaved` argument outright (`del interleaved`). That silently broke DeepSeek-V3 (`modeling_deepseek_v3.py`), which uses the same shared util with `interleaved=getattr(self.config, "rope_interleave", True)` and depends on the pairwise rotation. Move the fix to the right layer: - Restore the `interleaved=True` branch in the shared util (DSv3 contract). - Drop the `interleaved=getattr(self.config, "mrope_interleaved", False)` kwarg at the two Qwen3.5/3.6 attention call sites, so Qwen always uses half-rotate regardless of `mrope_interleaved`. - Document the convention split (Qwen half-rotate vs. DSv3 pairwise) at the shared util. - Restore pairwise-path tests and add an AST-based regression test that fails if either Qwen3.5/3.6 attention re-introduces an `interleaved=` kwarg on the rotary call. --- .../transformers/qwen3_5/modeling_qwen3_5.py | 11 +-- .../qwen3_5_moe/modeling_qwen3_5_moe.py | 11 +-- .../models/transformers/qwen3_5_shared.py | 6 ++ tests/models/test_qwen3_5_apply_rotary.py | 79 ++++++++++++++----- 4 files changed, 73 insertions(+), 34 deletions(-) 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/modeling_qwen3_5_moe.py b/src/xorl/models/transformers/qwen3_5_moe/modeling_qwen3_5_moe.py index 4ce0e1d9..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: diff --git a/src/xorl/models/transformers/qwen3_5_shared.py b/src/xorl/models/transformers/qwen3_5_shared.py index a547ff19..f042c695 100644 --- a/src/xorl/models/transformers/qwen3_5_shared.py +++ b/src/xorl/models/transformers/qwen3_5_shared.py @@ -51,6 +51,12 @@ 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 diff --git a/tests/models/test_qwen3_5_apply_rotary.py b/tests/models/test_qwen3_5_apply_rotary.py index e476c814..1182016b 100644 --- a/tests/models/test_qwen3_5_apply_rotary.py +++ b/tests/models/test_qwen3_5_apply_rotary.py @@ -1,7 +1,12 @@ +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 @@ -20,10 +25,21 @@ def _build_halved_cos_sin(batch: int, seq: int, head_dim: int) -> tuple[torch.Te return cos, sin -def _hf_reference_interleaved( +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]: - """HF-style reference for interleaved-layout q/k with halved cos/sin. + """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)`. @@ -46,7 +62,23 @@ def to_interleaved(x: torch.Tensor) -> torch.Tensor: return to_interleaved(q_embed_h), to_interleaved(k_embed_h) -def test_interleaved_matches_hf_reference(): +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) @@ -54,7 +86,7 @@ def test_interleaved_matches_hf_reference(): 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_interleaved(q, k, cos, sin) + 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) @@ -74,19 +106,26 @@ def test_interleaved_pairwise_rotation_d8(): torch.testing.assert_close(q_out, q, atol=1e-6, rtol=1e-6) -def test_non_interleaved_unchanged(): - """Non-interleaved path must be unaffected by the fix.""" - 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_out, k_out = qwen3_5_apply_rotary_pos_emb(q, k, cos, sin, interleaved=False) - - cos_u = cos.unsqueeze(2) - sin_u = sin.unsqueeze(2) - expected_q = q * cos_u + rotate_half(q) * sin_u - expected_k = k * cos_u + rotate_half(k) * sin_u - torch.testing.assert_close(q_out, expected_q, atol=1e-6, rtol=1e-6) - torch.testing.assert_close(k_out, expected_k, 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`." + ) From 60269e1bfb13fd7578fd24a3df1dcdbfae4ad2af Mon Sep 17 00:00:00 2001 From: Qingyang Wu Date: Tue, 26 May 2026 20:56:51 -0700 Subject: [PATCH 43/49] refactor: convert load-bearing asserts in arguments.py / data_loader.py to explicit errors * refactor: convert load-bearing asserts to explicit errors First batch toward. Converts the highest-stakes asserts in arguments.py (parallelism config validation) and data_loader.py (dataset shape and gradient-accumulation invariants) to explicit ValueError/RuntimeError/TypeError. These are runtime-load-bearing checks that python -O would silently strip if left as asserts. Other ~440 asserts across src/xorl/ are not touched in this PR; they need a per-call audit (runtime-check vs dev-invariant). Tracking the rest under. * style: collapse multi-line raise ValueError to satisfy ruff-format --- src/xorl/arguments.py | 36 ++++++++++++++++++++++-------------- src/xorl/data/data_loader.py | 23 ++++++++++++++--------- 2 files changed, 36 insertions(+), 23 deletions(-) diff --git a/src/xorl/arguments.py b/src/xorl/arguments.py index 18281bc4..e4868956 100644 --- a/src/xorl/arguments.py +++ b/src/xorl/arguments.py @@ -1072,9 +1072,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: @@ -1106,12 +1109,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( @@ -1152,13 +1157,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 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': ...}}, ...])." From 1549d96ec7300e99db8d36517d64980708104ddd Mon Sep 17 00:00:00 2001 From: Ashwinee Panda Date: Wed, 27 May 2026 10:10:39 -0700 Subject: [PATCH 44/49] Tune Muon Gram-NS chunk memory --- src/xorl/arguments.py | 8 ++++++++ src/xorl/optim/muon.py | 20 ++++++++++++++++---- src/xorl/optim/optimizer.py | 6 +++++- src/xorl/server/server_arguments.py | 8 ++++++++ tests/optim/test_muon.py | 11 ++++++++++- tests/server/test_server_arguments.py | 2 ++ tests/test_arguments.py | 2 ++ 7 files changed, 51 insertions(+), 6 deletions(-) diff --git a/src/xorl/arguments.py b/src/xorl/arguments.py index e4868956..42ba21b5 100644 --- a/src/xorl/arguments.py +++ b/src/xorl/arguments.py @@ -644,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={ @@ -690,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 diff --git a/src/xorl/optim/muon.py b/src/xorl/optim/muon.py index 1ecf4c51..e817d458 100644 --- a/src/xorl/optim/muon.py +++ b/src/xorl/optim/muon.py @@ -45,7 +45,7 @@ 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( @@ -187,6 +187,8 @@ 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 @@ -223,6 +225,7 @@ def __init__( 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", @@ -240,6 +243,8 @@ def __init__( 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 @@ -269,6 +274,7 @@ 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, @@ -460,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 @@ -524,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) @@ -549,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 diff --git a/src/xorl/optim/optimizer.py b/src/xorl/optim/optimizer.py index 5942f327..f09794e2 100644 --- a/src/xorl/optim/optimizer.py +++ b/src/xorl/optim/optimizer.py @@ -14,7 +14,7 @@ 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 @@ -209,6 +209,10 @@ def _get_optimizer_cls_and_kwargs( 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"), diff --git a/src/xorl/server/server_arguments.py b/src/xorl/server/server_arguments.py index 91a78b33..2c779bd6 100644 --- a/src/xorl/server/server_arguments.py +++ b/src/xorl/server/server_arguments.py @@ -361,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={ @@ -742,6 +749,7 @@ 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, diff --git a/tests/optim/test_muon.py b/tests/optim/test_muon.py index f3cf1215..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), @@ -398,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/server/test_server_arguments.py b/tests/server/test_server_arguments.py index 3d2f9b7b..c33491d0 100644 --- a/tests/server/test_server_arguments.py +++ b/tests/server/test_server_arguments.py @@ -234,6 +234,7 @@ def test_load_server_arguments_threads_muon_gram_newton_schulz_through_nested_co "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, @@ -263,6 +264,7 @@ def test_load_server_arguments_threads_muon_gram_newton_schulz_through_nested_co 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 diff --git a/tests/test_arguments.py b/tests/test_arguments.py index c973111e..6f68d29a 100644 --- a/tests/test_arguments.py +++ b/tests/test_arguments.py @@ -100,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, @@ -123,6 +124,7 @@ 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 From b9a79e11f0d5493e9e1d8cb682ef6a9c191e5bc6 Mon Sep 17 00:00:00 2001 From: Ashwinee Panda Date: Wed, 27 May 2026 18:02:31 -0700 Subject: [PATCH 45/49] fix(fsdp2): canonicalize FSDP-wrapped ReduceOp in BF16 a2a reduce-scatter MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(fsdp2): canonicalize FSDP-wrapped ReduceOp in BF16 a2a reduce-scatter FSDP emits a wrapped _ReduceOp (via dist._make_nccl_premul_sum(1.0)) when a custom gradient_divide_factor is set. The previous SUM/AVG check compared raw object identity and rejected the wrapped op, blocking configs that mix BF16StochasticAllToAllReduceScatter with custom divide factors. Add _canonical_reduce_op() that unwraps the inner RedOpType (SUM, AVG, PREMUL_SUM-as-SUM since the installer requires factor=1.0) and use it in both the validation check and the post-reduction divide branch. Add a CPU unit test that exercises raw RedOpType, wrapped ReduceOp, and the premul-sum factory. Not in the q35-397B winning recipe (R45 with this path was below default throughput) β€” ship as a robustness/compat fix. * review: tighten _canonical_reduce_op + mark CPU helper test - Strengthen PREMUL_SUM comment: call out the install-time gradient_divide_factor==1.0 check as load-bearing; relaxing it without inspecting the premul factor here would silently under-weight grads. - Drop the brittle str(op_type).rsplit('.', 1)[-1] fallback; unknown ops now fall through to the raw return and are rejected by the downstream SUM/AVG check. - Add @pytest.mark.cpu to test_canonical_reduce_op_accepts_fsdp_wrapped_ops so it's selected by the CPU CI lane (pytest -m cpu); previously only the module-level distributed mark applied, hiding it from the lane it was meant to guard. --- src/xorl/distributed/fsdp2/bf16_a2a_reduce.py | 33 +++++++++++++++++-- tests/distributed/test_bf16_a2a_reduce.py | 10 ++++++ 2 files changed, 40 insertions(+), 3 deletions(-) diff --git a/src/xorl/distributed/fsdp2/bf16_a2a_reduce.py b/src/xorl/distributed/fsdp2/bf16_a2a_reduce.py index fb3fe979..fb9d52c5 100644 --- a/src/xorl/distributed/fsdp2/bf16_a2a_reduce.py +++ b/src/xorl/distributed/fsdp2/bf16_a2a_reduce.py @@ -43,6 +43,30 @@ 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.""" @@ -65,8 +89,11 @@ def __call__( ) -> Optional[dist.Work]: if async_op: raise NotImplementedError("BF16StochasticAllToAllReduceScatter does not support async_op=True") - if op != ReduceOp.SUM and op != ReduceOp.AVG: - raise NotImplementedError(f"BF16StochasticAllToAllReduceScatter requires SUM or AVG op, got {op}") + 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 " @@ -87,7 +114,7 @@ def __call__( 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 == ReduceOp.AVG: + if op_type == ReduceOp.AVG: summed.div_(world_size) output_tensor.copy_(summed) return None diff --git a/tests/distributed/test_bf16_a2a_reduce.py b/tests/distributed/test_bf16_a2a_reduce.py index f3057b3e..7318c557 100644 --- a/tests/distributed/test_bf16_a2a_reduce.py +++ b/tests/distributed/test_bf16_a2a_reduce.py @@ -14,8 +14,10 @@ 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 @@ -29,6 +31,14 @@ 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"]) From 48efa243ccb5590da1f4193230ebc0c0a8529eb4 Mon Sep 17 00:00:00 2001 From: Qingyang Wu Date: Wed, 27 May 2026 18:36:00 -0700 Subject: [PATCH 46/49] chore(deps): bump flash-attn-cute to flash-attn-4 @ 59f01d6 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The previously-pinned commit (5678dd9, Feb 23, 2026) explicitly raises "varlen backward is not yet supported on sm90" β€” making FA4 unusable for packed-sequence training on H100. The newer commit (59f01d6, May 27, 2026) adds Hopper varlen backward support. Also renames the dep from flash-attn-cute to flash-attn-4 since upstream renamed the package between these commits. Verified on cu129 venv: FA4 varlen + mask_mod forward/backward all match torch SDPA reference within bf16 tolerance. Transitively bumps quack-kernels 0.3.10 -> 0.4.1. --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 5d1cc652..ef71612a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -40,7 +40,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", ] From c638722c98c6b52b0933e1f6feacfd49196ae138 Mon Sep 17 00:00:00 2001 From: Zhongzhu Zhou Date: Mon, 15 Jun 2026 23:36:33 -0700 Subject: [PATCH 47/49] feat(weight-sync): nccl_simple two-phase backend + sidecar routing for cross-node NCCL sync * feat(weight-sync): nccl_simple two-phase backend + sidecar routing for cross-node NCCL sync Adds a cross-node NCCL weight-sync path (trainer FSDP -> sglang TP) that avoids two deadlocks hit on 14B / FSDP=4 / 2-node: 1. nccl_simple backend (two-phase). nccl_broadcast interleaves the FSDP unshard all-gather (intra-node) and the weight-comm broadcast (inter-node) per module; with FSDP ranks 1..N racing ahead of rank 0, the two NCCL communicators enqueue kernels in different orders across ranks and deadlock (observed consistently at the 7th bucket). nccl_simple removes the interleaving: Phase A stages all params to CPU during the FSDP loop (FSDP collectives only), Phase B broadcasts in 1 GiB chunks (weight comm only). Select with --server.sync_inference_method=nccl_simple. 2. NCCL backend cache (XORL_NCCL_BACKEND_CACHE, default on). The NCCL comm init (TCPStore rendezvous + ncclCommInitRankConfig) must happen exactly once; building a new backend each sync re-sends /init_weights_update_group, making sglang attempt a second ncclCommInitRankConfig on the same group name and deadlocking the first comm. Cache the backend per (endpoint, group_name, world_size); initialize() runs once and is reused across syncs. 3. Sidecar-receiver routing (XORL_WEIGHT_SYNC_PORT). Receiving the broadcast inside sglang's NUMA-pinned scheduler starves torch's nonblocking eager NCCL init (device_id=) -- NCCL logs "Init START" but never "Init COMPLETE", while the broadcast kernel never launches. The identical receive in fresh processes runs at ~24 GB/s. This env overrides the receiver port for init/update/destroy so they can target a fresh-process sidecar receiver, while generation and pause/resume keep hitting the real serving port. Also: - HTTP fail-fast + per-endpoint thread logging in nccl_broadcast: if the update HTTP call errors immediately (connection refused), raise before dist.broadcast instead of blocking forever on receivers that were never notified (no more silent 600s hangs). - pause_mode default retract -> in_place (api_types / operations / remote): retract blocked sglang's event loop during the sync. Validated 2-node (trainer FSDP=4 -> sglang TP=4, Qwen2.5-Coder-14B): 31.10 GB / 580 params / 51 buckets synced, ~40s cold (incl. NCCL group init) / ~15s steady-state, repeatable across training steps. Touches weight sync -> run the K3 logprob comparison before merge. * style: apply ruff lint + format fixes --------- Co-authored-by: Ashwinee Panda --- src/xorl/server/api_server/api_types.py | 2 +- src/xorl/server/backend/remote.py | 2 +- src/xorl/server/protocol/operations.py | 2 +- src/xorl/server/server_arguments.py | 7 +- .../server/weight_sync/backends/__init__.py | 6 +- .../weight_sync/backends/nccl_broadcast.py | 47 ++++++- .../weight_sync/backends/nccl_simple.py | 123 ++++++++++++++++++ src/xorl/server/weight_sync/handler.py | 65 ++++++++- 8 files changed, 237 insertions(+), 17 deletions(-) create mode 100644 src/xorl/server/weight_sync/backends/nccl_simple.py diff --git a/src/xorl/server/api_server/api_types.py b/src/xorl/server/api_server/api_types.py index 7915eadf..6190dfb7 100644 --- a/src/xorl/server/api_server/api_types.py +++ b/src/xorl/server/api_server/api_types.py @@ -903,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. " diff --git a/src/xorl/server/backend/remote.py b/src/xorl/server/backend/remote.py index f1edab96..f1358687 100644 --- a/src/xorl/server/backend/remote.py +++ b/src/xorl/server/backend/remote.py @@ -355,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, diff --git a/src/xorl/server/protocol/operations.py b/src/xorl/server/protocol/operations.py index 309a5644..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 diff --git a/src/xorl/server/server_arguments.py b/src/xorl/server/server_arguments.py index 2c779bd6..6f76d583 100644 --- a/src/xorl/server/server_arguments.py +++ b/src/xorl/server/server_arguments.py @@ -584,11 +584,14 @@ class ServerArguments: # Inference Weight Sync Configuration # ======================================================================== - sync_inference_method: Literal["nccl_broadcast", "p2p"] = 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)" }, diff --git a/src/xorl/server/weight_sync/backends/__init__.py b/src/xorl/server/weight_sync/backends/__init__.py index 32de8fc8..b3f49502 100644 --- a/src/xorl/server/weight_sync/backends/__init__.py +++ b/src/xorl/server/weight_sync/backends/__init__.py @@ -29,8 +29,12 @@ def create_backend( from .nccl_broadcast import NCCLBroadcastBackend # noqa: PLC0415 return NCCLBroadcastBackend(config, **kwargs) + 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', 'p2p'.") + raise ValueError(f"Unknown weight sync backend: {method!r}. Supported: 'nccl_broadcast', 'nccl_simple', 'p2p'.") diff --git a/src/xorl/server/weight_sync/backends/nccl_broadcast.py b/src/xorl/server/weight_sync/backends/nccl_broadcast.py index c06929e9..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.""" @@ -236,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) @@ -399,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, @@ -455,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) @@ -540,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]: @@ -650,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, @@ -673,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) @@ -682,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) @@ -703,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() 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/handler.py b/src/xorl/server/weight_sync/handler.py index 6e2cc9b3..8160140b 100644 --- a/src/xorl/server/weight_sync/handler.py +++ b/src/xorl/server/weight_sync/handler.py @@ -214,9 +214,20 @@ def _safe_abort_token(value: Optional[str]) -> str: _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 + 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) @@ -224,6 +235,13 @@ def _atexit_destroy_cached_backend() -> None: 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) @@ -767,8 +785,13 @@ def _sync_weights( # 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 - cache_enabled = os.environ.get("XORL_P2P_BACKEND_CACHE", "1") == "1" and sync_method == "p2p" + 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 = ( @@ -790,8 +813,14 @@ def _sync_weights( ) ), ) - if ( - cache_enabled + + # 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) @@ -811,6 +840,13 @@ def _sync_weights( 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 @@ -826,7 +862,16 @@ def _sync_weights( 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(): @@ -836,6 +881,14 @@ def _sync_weights( } 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: From bb55abb6e92e175de9a3f1eb3f93643abd315f7f Mon Sep 17 00:00:00 2001 From: Ashwinee Panda Date: Mon, 15 Jun 2026 23:39:21 -0700 Subject: [PATCH 48/49] fix(server): unblock multi-LoRA e2e on main (forward-path fixes + Qwen3-30B-A3B config) * fix(server): keep forward evals out of inference_mode * Add Qwen3-30B-A3B routed-expert multi-LoRA server config * fix(server): use returned metric ops in forward loop --- .../configs/lora/qwen3_30b_a3b_lora.yaml | 75 +++++++++++++++++++ src/xorl/server/runner/model_runner.py | 10 +-- tests/server/runner/test_opd_runner.py | 72 ++++++++++++++++++ 3 files changed, 152 insertions(+), 5 deletions(-) create mode 100644 examples/server/configs/lora/qwen3_30b_a3b_lora.yaml 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/src/xorl/server/runner/model_runner.py b/src/xorl/server/runner/model_runner.py index ebb9828b..e5b28236 100644 --- a/src/xorl/server/runner/model_runner.py +++ b/src/xorl/server/runner/model_runner.py @@ -2266,7 +2266,7 @@ def _profile_phase_elapsed(start: float) -> float: # 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, metric_ops) + self._accumulate_loss_metrics(accumulated_loss_metrics, is_metrics, loss_fn, is_metric_ops) # Cleanup del micro_batch, outputs, local_loss_sum @@ -2680,10 +2680,10 @@ def forward_backward( return result - # inference_mode (vs no_grad) skips view tracking and version counters, - # which is safe here because the returned tensors are only consumed by - # logging / cross-rank reductions that never re-enter autograd. - @torch.inference_mode() + # 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]], diff --git a/tests/server/runner/test_opd_runner.py b/tests/server/runner/test_opd_runner.py index ea2d4c87..9b9d3c31 100644 --- a/tests/server/runner/test_opd_runner.py +++ b/tests/server/runner/test_opd_runner.py @@ -50,6 +50,37 @@ def forward(self, input): 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) @@ -117,6 +148,47 @@ def fake_all_reduce(_tensor, op=None, group=None): 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) From a29857ced333433de3817d606b71f0fa66b0aa19 Mon Sep 17 00:00:00 2001 From: Ashwinee Panda Date: Tue, 16 Jun 2026 01:37:39 -0700 Subject: [PATCH 49/49] =?UTF-8?q?fix:=20un-vendor=20FLA=20=E2=80=94=20use?= =?UTF-8?q?=20upstream=20flash-linear-attention=20for=20GatedDeltaNet?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace the stale vendored flash-linear-attention fork with upstream FLA for the GatedDeltaNet kernel + norms, fixing the silent correctness bug. - GDN chunk/recurrent kernels -> fla.ops.gated_delta_rule. The vendored fork ran the Triton gated chunk_bwd_dqkwg, which produces WRONG gradients on Triton>=3.4 / Hopper with no guard; upstream routes that backward to a TileLang kernel. - Norms (RMSNorm / FusedRMSNormGated) -> fla.modules (numerically identical, cos=1.0 fwd+bwd). - Deps: flash-linear-attention pinned to public commit fla-org@97bcb883; tilelang==0.1.11 (first PyPI release with both the b_dq layout fix and Hopper fp32-MMA support needed by the Ulysses CP backward); apache-tvm-ffi==0.1.11. - Kept xorl's pure-torch ShortConvolution + ops/cp CP-halo glue on purpose: upstream's conv is correct but ~3x slower in GDN training (per-call host overhead x ~180 conv calls/step), measured via a same-node A/B. Validated: GPU CI 427 passed; CP-equivalence + Ulysses smoke green; end-to-end loss-match on Qwen3.5-35B-A3B @ 65k. --- pyproject.toml | 5 + src/xorl/ops/linear_attention/__init__.py | 3 +- .../linear_attention/layers/gated_deltanet.py | 13 +- .../ops/linear_attention/modules/__init__.py | 4 - .../modules/fused_norm_gate.py | 89 -- .../ops/linear_attention/modules/l2norm.py | 266 ---- .../ops/linear_attention/modules/layernorm.py | 57 - src/xorl/ops/linear_attention/ops/__init__.py | 2 +- .../ops/common/chunk_delta_h.py | 597 -------- .../linear_attention/ops/common/chunk_o.py | 707 ---------- .../ops/common/chunk_scaled_dot_kkt.py | 127 -- .../linear_attention/ops/cp/chunk_delta_h.py | 1233 ----------------- .../ops/gated_delta_rule/__init__.py | 8 - .../ops/gated_delta_rule/chunk.py | 383 ----- .../ops/gated_delta_rule/fused_recurrent.py | 249 ---- .../ops/gated_delta_rule/wy_fast.py | 345 ----- .../linear_attention/ops/utils/__init__.py | 15 - .../ops/linear_attention/ops/utils/cumsum.py | 474 ------- .../ops/linear_attention/ops/utils/index.py | 38 - src/xorl/ops/linear_attention/ops/utils/op.py | 55 - .../linear_attention/ops/utils/solve_tril.py | 399 ------ tests/ops/test_gated_delta_rule.py | 3 +- 22 files changed, 17 insertions(+), 5055 deletions(-) delete mode 100644 src/xorl/ops/linear_attention/modules/fused_norm_gate.py delete mode 100644 src/xorl/ops/linear_attention/modules/l2norm.py delete mode 100644 src/xorl/ops/linear_attention/modules/layernorm.py delete mode 100644 src/xorl/ops/linear_attention/ops/common/chunk_delta_h.py delete mode 100644 src/xorl/ops/linear_attention/ops/common/chunk_o.py delete mode 100644 src/xorl/ops/linear_attention/ops/common/chunk_scaled_dot_kkt.py delete mode 100644 src/xorl/ops/linear_attention/ops/cp/chunk_delta_h.py delete mode 100644 src/xorl/ops/linear_attention/ops/gated_delta_rule/__init__.py delete mode 100644 src/xorl/ops/linear_attention/ops/gated_delta_rule/chunk.py delete mode 100644 src/xorl/ops/linear_attention/ops/gated_delta_rule/fused_recurrent.py delete mode 100644 src/xorl/ops/linear_attention/ops/gated_delta_rule/wy_fast.py delete mode 100644 src/xorl/ops/linear_attention/ops/utils/__init__.py delete mode 100644 src/xorl/ops/linear_attention/ops/utils/cumsum.py delete mode 100644 src/xorl/ops/linear_attention/ops/utils/index.py delete mode 100644 src/xorl/ops/linear_attention/ops/utils/op.py delete mode 100644 src/xorl/ops/linear_attention/ops/utils/solve_tril.py diff --git a/pyproject.toml b/pyproject.toml index ef71612a..4d134f14 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -24,6 +24,11 @@ dependencies = [ "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", 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/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 = [