From 74de9e3fa15604b2046b10838f7ff6a868fd4fcc Mon Sep 17 00:00:00 2001 From: Qingyang Wu Date: Wed, 19 Aug 2026 04:35:28 +0000 Subject: [PATCH 1/2] Add GLM-5.2 adapted-target scope and path-qualified LoRA targets GLM-5.2's block-FP8 QLoRA lane builds a complete deterministic inventory and rejects `lora_target_modules`, so there was no way to isolate where MoE learning happens. Add `glm52_lora_scope` (all | moe | shared_experts | routed_experts). Scope selects which factors TRAIN, not which modules are adapted. The complete exact family is always constructed because NativeBlockFP8Linear implements no activation backward ("phase-one forward is scoring-only"), so a region left unadapted blocks gradients from reaching adapted regions downstream of it. Out-of-scope factors are frozen with lora_B == 0 and therefore contribute nothing to the forward: step one still reproduces the frozen-base loss exactly, and the forward program is identical for every scope. Three checks assumed "adapted implies trainable", which held while the complete family was the only supported configuration: * the exact LM head required both factor masters to be trainable; the real invariant is that the pair SHARES trainability, since a half-frozen pair would take a gradient on one side of a product whose other side is fixed * adapter gradient ownership required a gradient, a staged gradient, and a compiled-producer numerator for every inventory factor; a frozen factor produces none by design, and absence is still a fault for factors that are supposed to be learning Separately, generalize target selection for every architecture: a target containing a path separator or glob is matched with fnmatch against the full module path. Leaf-name matching discards a module's position in the tree, so "gate_proj" cannot distinguish a routed expert from a shared expert or a dense MLP. Bare names keep their existing leaf/indirect semantics, and patterns join the same fail-closed coverage check. Also decorate the worker entrypoint with torch elastic's `record` and enable faulthandler, so a rank that dies mid-forward reports a traceback instead of `error_file: ` and a bare exit code. Measured on 16x H100 (rank 64, 64 steps, lr 1e-4), password memorization: full scope (1,700 trainable factors) reaches 0.0328; routed experts only (450 factors) reaches 0.0645 -- about half the improvement on 26% of the factors. Note: the exported adapter is the full inventory regardless of scope, with untrained factors stored as zeros. --- docs/k3/LORA_CONTRACT.md | 15 ++ .../configs/qlora/glm5_2_qlora_block_fp8.yaml | 112 ++++++++++ .../qlora/glm5_2_qlora_routed_experts.yaml | 121 ++++++++++ .../password_memorization/GLM52_RUNBOOK.md | 124 ++++++++++ .../run_glm52_password_train.py | 211 ++++++++++++++++++ src/xorl/lora/utils.py | 29 ++- src/xorl/models/auto.py | 12 +- .../transformers/glm5/exact_lm_head_qlora.py | 12 +- src/xorl/models/transformers/glm5/qlora.py | 120 +++++++--- src/xorl/server/runner/adapters/manager.py | 23 +- src/xorl/server/runner/model_runner.py | 1 + src/xorl/server/runner/setup.py | 26 ++- src/xorl/server/server_arguments.py | 34 ++- src/xorl/trainers/model_builder.py | 2 + tests/models/test_glm52_lora_scope.py | 142 ++++++++++++ .../models/test_lora_path_pattern_targets.py | 108 +++++++++ 16 files changed, 1055 insertions(+), 37 deletions(-) create mode 100644 examples/server/configs/qlora/glm5_2_qlora_block_fp8.yaml create mode 100644 examples/server/configs/qlora/glm5_2_qlora_routed_experts.yaml create mode 100644 examples/server/password_memorization/GLM52_RUNBOOK.md create mode 100644 examples/server/password_memorization/run_glm52_password_train.py create mode 100644 tests/models/test_glm52_lora_scope.py create mode 100644 tests/models/test_lora_path_pattern_targets.py diff --git a/docs/k3/LORA_CONTRACT.md b/docs/k3/LORA_CONTRACT.md index 153fbe2f..e6096acd 100644 --- a/docs/k3/LORA_CONTRACT.md +++ b/docs/k3/LORA_CONTRACT.md @@ -57,6 +57,21 @@ Trainer-only autograd supplies activation and factor gradients from the effective BF16 factor values. It is a checked straight-through treatment of the frozen quantized base, not a derivative of FP8 quantization. +## Adapted-target scope (GLM-5.2) + +`glm52_lora_scope` (`all` | `moe` | `shared_experts` | `routed_experts`) selects +which factors **train**, never which modules are adapted. The complete family is +always constructed: `NativeBlockFP8Linear` implements no activation backward, so +a region left unadapted would block gradients from reaching adapted regions +downstream of it. Out-of-scope factors are frozen with `lora_B == 0`, so the +forward program is identical for every scope and step one reproduces the +frozen-base loss exactly. + +Only `all` at rank 1, alpha 1 is the qualified configuration. Narrowed scopes +run the same forward but train a different parameter set, so they carry no +bit-exactness claim; the exported adapter still contains the full inventory, +with untrained factors stored as zeros. + ## Unsupported inheritance Neither lane implicitly covers multiple simultaneously active adapters, diff --git a/examples/server/configs/qlora/glm5_2_qlora_block_fp8.yaml b/examples/server/configs/qlora/glm5_2_qlora_block_fp8.yaml new file mode 100644 index 00000000..06495c44 --- /dev/null +++ b/examples/server/configs/qlora/glm5_2_qlora_block_fp8.yaml @@ -0,0 +1,112 @@ +# Server-side configuration for XORL Training Server +# GLM-5.2 exact active-LoRA over the official native block-FP8 base. +# +# 16x H100 80GB (WORLD16/PP1/TP1/DP1/EP16/CP16) — see docs/k3/DEFAULTS_AND_PARETO.md. +# The official FP8 checkpoint is ~704 GB on disk; EP16 puts ~46 GB of expert +# weight per rank before activations. This does NOT fit on fewer GPUs. +# +# Unlike every Qwen LoRA/QLoRA config in this directory, the adapted target set +# is NOT configurable here. GLM-5.2 block-FP8 QLoRA builds its complete +# deterministic inventory (625 adapted linears + 75 routed expert banks = +# 1,700 FP32 factors) and rejects `lora_target_modules` / +# `lora_target_manifest` / `qlora_exclude_modules` outright. +# See src/xorl/models/transformers/glm5/qlora.py and docs/k3/LORA_CONTRACT.md. + +# ============================================================================ +# Model Configuration +# ============================================================================ +model_path: zai-org/GLM-5.2-FP8 +tokenizer_path: zai-org/GLM-5.2-FP8 +# Canonical GLM-5.2 is a numerical PROGRAM, not a set of options: the exact +# forward requires flash_attention_4 and rejects flash_attention_3 outright +# (src/xorl/models/auto.py:722). Leaving it unset would also resolve correctly. +attn_implementation: flash_attention_4 +moe_implementation: triton # required by the GLM-5.2 block-FP8 QLoRA admission +ep_dispatch: alltoall # alltoall selects the EXACT active-LoRA family + # (deepep would build the ordinary block-FP8 lane) + +# ============================================================================ +# Parallelism Configuration +# ============================================================================ +data_parallel_mode: fsdp2 +tensor_parallel_size: 1 # adapter-bearing modules require body TP1 +lm_head_tensor_parallel_size: 16 # exact lm-head component requires TP16 +expert_parallel_size: 16 # exact routed-expert component requires EP16 +ulysses_parallel_size: 16 +ringattn_parallel_size: 1 # GLM-5 DSA does not support ring attention +data_parallel_replicate_size: 1 +data_parallel_shard_size: 1 # LoRA + EP asserts ep_fsdp_size == 1 +# The exact GLM-5.2 lm head is a separate FSDP unit computing vocab-parallel CE +# against its local shard; torch_parallelize.py:563 refuses to build without it. +fsdp_sharded_lm_head_loss: true +# lm_head_tensor_parallel_size>1 requires cp_fsdp_mode='all' so the external SP +# gradient sync does not double-reduce against the lm_head replica reduction +# (torch_parallelize.py:594). This is the default; set explicitly because the +# exact lane depends on it. +cp_fsdp_mode: all + +# ============================================================================ +# Memory & Performance +# ============================================================================ +enable_mixed_precision: true +enable_gradient_checkpointing: true +enable_full_shard: true +enable_activation_offload: false +init_device: meta +load_weights_mode: all_ranks + +# ============================================================================ +# Checkpointing +# ============================================================================ +output_dir: outputs/GLM-5.2-server-qlora-password +load_checkpoint_path: "" +ckpt_manager: dcp + +# ============================================================================ +# Logging +# ============================================================================ +log_level: INFO + +# ============================================================================ +# Worker Configuration +# ============================================================================ +worker_connection_timeout: 180.0 +worker_max_retries: 5 + +# ============================================================================ +# Data Processing Configuration +# ============================================================================ +sample_packing_sequence_len: 32768 +enable_packing: true + +# ============================================================================ +# QLoRA Configuration (GLM-5.2 exact active-LoRA) +# ============================================================================ +enable_lora: true +enable_qlora: true +block_fp8_qlora_training: true +quant_format: block_fp8 +quant_group_size: 128 +moe_hybrid_shared_lora: true +lora_export_format: sglang_shared_outer +freeze_router: true +merge_qkv: true +# NOTE: docs/k3/LORA_CONTRACT.md qualifies ONLY rank 1 / alpha 1 for the exact +# lane. Rank 64 is accepted by the code (glm52_exact_lora_scaling requires only +# positive integers) but is OUTSIDE the qualified configuration, so the +# train/serve bit-exactness claim does not cover it. +# alpha == rank keeps scaling at 1.0, matching the rank-1 baseline. +# The exact lane requires max_lora_rank == lora_rank (defaults to lora_rank). +lora_rank: 64 +lora_alpha: 64 +# NOTE: lora_target_modules is intentionally absent — setting it is rejected. + +skip_initial_checkpoint: true + +# ============================================================================ +# Deliberately omitted — canonical GLM-5.2 resolves these itself +# (docs/src/content/docs/config-reference/server.md, "Numerical alignment flags"): +# sparse_mla_enabled -> true, sparse_mla_backend -> flashmla, +# rmsnorm_mode -> sglang_fused, rope_native -> true, rope_class_b -> true, +# attention_cast_bf16 -> false, ce_mode -> bi_fused. +# Setting them by hand risks selecting a different numerical program. diff --git a/examples/server/configs/qlora/glm5_2_qlora_routed_experts.yaml b/examples/server/configs/qlora/glm5_2_qlora_routed_experts.yaml new file mode 100644 index 00000000..01b504d5 --- /dev/null +++ b/examples/server/configs/qlora/glm5_2_qlora_routed_experts.yaml @@ -0,0 +1,121 @@ +# Server-side configuration for XORL Training Server +# GLM-5.2 exact active-LoRA over the official native block-FP8 base. +# +# 16x H100 80GB (WORLD16/PP1/TP1/DP1/EP16/CP16) — see docs/k3/DEFAULTS_AND_PARETO.md. +# The official FP8 checkpoint is ~704 GB on disk; EP16 puts ~46 GB of expert +# weight per rank before activations. This does NOT fit on fewer GPUs. +# +# Unlike every Qwen LoRA/QLoRA config in this directory, the adapted target set +# is NOT configurable here. GLM-5.2 block-FP8 QLoRA builds its complete +# deterministic inventory (625 adapted linears + 75 routed expert banks = +# 1,700 FP32 factors) and rejects `lora_target_modules` / +# `lora_target_manifest` / `qlora_exclude_modules` outright. +# See src/xorl/models/transformers/glm5/qlora.py and docs/k3/LORA_CONTRACT.md. + +# ============================================================================ +# Model Configuration +# ============================================================================ +model_path: zai-org/GLM-5.2-FP8 +tokenizer_path: zai-org/GLM-5.2-FP8 +# Canonical GLM-5.2 is a numerical PROGRAM, not a set of options: the exact +# forward requires flash_attention_4 and rejects flash_attention_3 outright +# (src/xorl/models/auto.py:722). Leaving it unset would also resolve correctly. +attn_implementation: flash_attention_4 +moe_implementation: triton # required by the GLM-5.2 block-FP8 QLoRA admission +# A scoped run keeps the exact FORWARD program (Class-B RoPE, sglang_fused +# RMSNorm, sparse MLA, canonical packed_ep16_v2 MoE transport) and therefore the +# exact lane's transport too. Only the adapted target set is narrower, so results +# stay comparable with a full-scope run. +ep_dispatch: alltoall + +# ============================================================================ +# Parallelism Configuration +# ============================================================================ +data_parallel_mode: fsdp2 +tensor_parallel_size: 1 # adapter-bearing modules require body TP1 +lm_head_tensor_parallel_size: 16 # exact lm-head component requires TP16 +expert_parallel_size: 16 # exact routed-expert component requires EP16 +ulysses_parallel_size: 16 +ringattn_parallel_size: 1 # GLM-5 DSA does not support ring attention +data_parallel_replicate_size: 1 +data_parallel_shard_size: 1 # LoRA + EP asserts ep_fsdp_size == 1 +# The exact GLM-5.2 lm head is a separate FSDP unit computing vocab-parallel CE +# against its local shard; torch_parallelize.py:563 refuses to build without it. +fsdp_sharded_lm_head_loss: true +# lm_head_tensor_parallel_size>1 requires cp_fsdp_mode='all' so the external SP +# gradient sync does not double-reduce against the lm_head replica reduction +# (torch_parallelize.py:594). This is the default; set explicitly because the +# exact lane depends on it. +cp_fsdp_mode: all + +# ============================================================================ +# Memory & Performance +# ============================================================================ +enable_mixed_precision: true +enable_gradient_checkpointing: true +enable_full_shard: true +enable_activation_offload: false +init_device: meta +load_weights_mode: all_ranks + +# ============================================================================ +# Checkpointing +# ============================================================================ +output_dir: outputs/GLM-5.2-server-qlora-routed +load_checkpoint_path: "" +ckpt_manager: dcp + +# ============================================================================ +# Logging +# ============================================================================ +log_level: INFO + +# ============================================================================ +# Worker Configuration +# ============================================================================ +worker_connection_timeout: 180.0 +worker_max_retries: 5 + +# ============================================================================ +# Data Processing Configuration +# ============================================================================ +sample_packing_sequence_len: 32768 +enable_packing: true + +# ============================================================================ +# QLoRA Configuration (GLM-5.2 exact active-LoRA) +# ============================================================================ +# Adapt ONLY the 75 routed-expert banks (450 factors). Attention, dense MLPs, +# shared experts, and the lm_head stay frozen native FP8. This isolates how much +# of the task is learned in the routed experts alone -- the closest GLM-5.2 +# analogue of the Qwen reference recipe's MoE-only adapters. +glm52_lora_scope: routed_experts + +enable_lora: true +enable_qlora: true +block_fp8_qlora_training: true +quant_format: block_fp8 +quant_group_size: 128 +moe_hybrid_shared_lora: true +lora_export_format: sglang_shared_outer +freeze_router: true +merge_qkv: true +# NOTE: docs/k3/LORA_CONTRACT.md qualifies ONLY rank 1 / alpha 1 for the exact +# lane. Rank 64 is accepted by the code (glm52_exact_lora_scaling requires only +# positive integers) but is OUTSIDE the qualified configuration, so the +# train/serve bit-exactness claim does not cover it. +# alpha == rank keeps scaling at 1.0, matching the rank-1 baseline. +# The exact lane requires max_lora_rank == lora_rank (defaults to lora_rank). +lora_rank: 64 +lora_alpha: 64 +# NOTE: lora_target_modules is intentionally absent — setting it is rejected. + +skip_initial_checkpoint: true + +# ============================================================================ +# Deliberately omitted — canonical GLM-5.2 resolves these itself +# (docs/src/content/docs/config-reference/server.md, "Numerical alignment flags"): +# sparse_mla_enabled -> true, sparse_mla_backend -> flashmla, +# rmsnorm_mode -> sglang_fused, rope_native -> true, rope_class_b -> true, +# attention_cast_bf16 -> false, ce_mode -> bi_fused. +# Setting them by hand risks selecting a different numerical program. diff --git a/examples/server/password_memorization/GLM52_RUNBOOK.md b/examples/server/password_memorization/GLM52_RUNBOOK.md new file mode 100644 index 00000000..fdc0dac5 --- /dev/null +++ b/examples/server/password_memorization/GLM52_RUNBOOK.md @@ -0,0 +1,124 @@ +# GLM-5.2 password-memorization LoRA — runbook + +Adapts the recipe behind +[`togethercomputer/Qwen3-30B-A3B-MoE-LoRA-Password-Adapters`](https://huggingface.co/togethercomputer/Qwen3-30B-A3B-MoE-LoRA-Password-Adapters) +to GLM-5.2. Trained adapters: +[`togethercomputer/GLM-5.2-Password-LoRA-xorl`](https://huggingface.co/togethercomputer/GLM-5.2-Password-LoRA-xorl). + +## Results + +| Trainable factors | Steps | LR | Final loss | +|---|---:|---:|---:| +| 1,700 (`glm52_lora_scope: all`) | 64 | 1e-4 | **0.0328** | +| 450 (`glm52_lora_scope: routed_experts`) | 64 | 1e-4 | **0.0645** | + +Rank 64 / alpha 64, 8-step warmup + cosine, ~500 label tokens/step, 16x H100, +~10.7 s/step. The two runs differ **only** in which factors were trainable, so +the pair measures how much of the task lives in the routed experts: they reach +about half the improvement on 26% of the factors. + +Recall has **not** been verified by generation. That needs a GLM-5.2 SGLang +endpoint and a weight sync (Phase 2 below). The reference card verified by +teacher-forced generation, which is the test that actually settles this. + +## Two things that cost hours — read first + +**1. RDMA must be requested explicitly.** Without `rdma/infiniband` in the pod +resources (plus `IPC_LOCK`), `/dev/infiniband` is absent, NCCL silently falls +back to TCP, and steps take **~1,070 s instead of ~9 s** — a ~120x penalty that +looks like a hang, not a misconfiguration. See `~/k8s-setup/glm52-train-16gpu.yaml`. + +**2. Client timeouts must exceed server timeouts.** The driver's per-future wait +must be above the server's `--operation-timeout`, or the client reports a "hang" +long before the server's own verdict arrives. + +## What differs from the Qwen recipe + +The Qwen adapters were MoE-only by *excluding* attention from +`lora_target_modules`. GLM-5.2 rejects that field entirely and builds a complete +deterministic inventory: **1,700 factor tensors** over attention (390 targets), +routed experts (75 banks), shared experts (225), dense MLPs (9) and `lm_head`. +Router and DSA indexer stay frozen. + +Isolation is expressed with **`glm52_lora_scope`** (`all` | `moe` | +`shared_experts` | `routed_experts`), which selects which factors **train**, not +which modules are adapted: + +```yaml +glm52_lora_scope: routed_experts # 450 of 1,700 factors train +ep_dispatch: alltoall # required: the exact family is still built +``` + +Every region keeps its exact adapter module because `NativeBlockFP8Linear` is +**forward-only** ("phase-one forward is scoring-only; activation backward +requires a validated kernel"). A region left unadapted blocks gradients from +reaching adapted regions downstream of it. Out-of-scope factors are frozen with +`lora_B == 0`, so they contribute nothing to the forward and step 1 reproduces +the frozen-base loss exactly. + +Consequence: **the exported adapter is the full inventory regardless of scope** +(16 GB at rank 64), with untrained factors stored as zeros. + +Only `scope: all` at rank 1 / alpha 1 is qualified for train/serve bit-exactness +(`docs/k3/LORA_CONTRACT.md`). Rank 64 and narrowed scopes run the same forward +program but carry no such claim. + +## Prerequisites + +The exact lane imports `sglang.srt.*` and `sglang.kernels.ops.gemm.*` on every +adapted forward, so it needs the **combined torch-2.11 environment**, not the +default profile: + +```bash +git submodule update --init --recursive +cp pyproject.sglang.toml pyproject.toml # restore the original afterwards +UV_PROJECT_ENVIRONMENT=.venv-sglang uv sync +``` + +Weights: `zai-org/GLM-5.2-FP8`, 141 shards / ~704 GB. Topology is fixed at +WORLD16 / PP1 / TP1 / DP1 / EP16 / CP16 with lm-head TP16; it does not fit in +fewer GPUs. + +## Phase 1 — training (16 GPUs) + +```bash +kubectl apply -f ~/k8s-setup/glm52-train-16gpu.yaml +kubectl logs -f glm52-train-0 -n qywu # ~6 min to load 141 shards + +python examples/server/password_memorization/run_glm52_password_train.py \ + --model zai-org/GLM-5.2-FP8 \ + --train-url http://:6000 \ + --steps 64 --lr 1e-4 --lr-schedule warmup_cosine --warmup-steps 8 \ + --repeat 16 --model-id my-run --save-name my-adapter +``` + +Sanity checks, in order: + +* `Registered adapter ... num_params=1700` — the complete family is built +* `GLM-5.2 LoRA scope 'routed_experts': froze 1250 of 1700 factors` — scope applied +* `Step 1/64: loss=2.2714929580688477` — frozen-base loss, so factors are fresh + +`--repeat N` replicates the 3 examples per step; at `--repeat 1` the batch packs +to ~128 tokens, which is 8 per CP rank and leaves nearly all 256 routed experts +empty. Each run needs its own `--model-id`: the reserved `default` session +cannot be unloaded, so reusing it silently inherits the previous adapter (the +driver aborts if step 1 is not the frozen-base loss). + +Steps take ~9-11 s. Anything near 1,000 s means NCCL is on TCP — check +`/dev/infiniband` inside the pod. + +## Phase 2 — recall verification (a second 16 GPUs) + +Not yet run. `run_password_test.py` syncs weights to SGLang and queries recall; +it needs a GLM-5.2 serving deployment, which does not fit alongside the trainer. + +## Known gaps + +* Recall unverified by generation (above). +* Narrowed scopes export the full inventory; filtering zero factors would shrink + the artifact substantially. +* No CPU-level construction test. Three bugs in the scope feature + (`NativeBlockFP8Linear` backward, the lm-head trainability assertion, and + gradient-ownership presence checks) were each found only by a ~7-minute load on + 16 GPUs. A miniature GLM-5.2 config exercising `_validate_constructed_model` + would catch that class on CPU in seconds. diff --git a/examples/server/password_memorization/run_glm52_password_train.py b/examples/server/password_memorization/run_glm52_password_train.py new file mode 100644 index 00000000..0d0d1b53 --- /dev/null +++ b/examples/server/password_memorization/run_glm52_password_train.py @@ -0,0 +1,211 @@ +#!/usr/bin/env python3 +"""GLM-5.2 password memorization — training half only. + +`run_password_test.py` requires a live SGLang endpoint because it verifies +recall after a weight sync. Serving GLM-5.2 costs a second 16-GPU allocation, +so this driver stops after training: it proves the exact active-LoRA lane +builds, takes gradients, and drives the loss to ~0 on the same task the +Qwen3-30B-A3B password adapters were trained on +(togethercomputer/Qwen3-30B-A3B-MoE-LoRA-Password-Adapters). + +Everything except the inference half is reused from run_password_test.py. + +What deliberately differs from the Qwen recipe: + * Target modules are NOT selectable. The Qwen adapters were MoE-only + (gate/up/down, no attention); GLM-5.2 builds its complete deterministic + inventory or refuses to start. See docs/k3/LORA_CONTRACT.md. + * Rank/alpha come from the server YAML, not from here. The qualified GLM-5.2 + configuration is rank 1 / alpha 1, against the Qwen recipe's rank 16. + +Usage (against the server from glm5_2_qlora_block_fp8.yaml): + + python run_glm52_password_train.py \ + --model zai-org/GLM-5.2-FP8 \ + --train-url http://glm52-train-0.glm52-train.qywu.svc.cluster.local:6000 \ + --steps 64 --lr 5e-4 --lr-schedule warmup_cosine --warmup-steps 8 +""" + +import argparse +import sys +import time +from pathlib import Path + +import requests +from transformers import AutoTokenizer + + +sys.path.insert(0, str(Path(__file__).resolve().parent)) + +from run_password_test import ( # noqa: E402 + CODES, + MODEL_ID, + _raise_on_failed_future, + build_training_data, + create_model, + get_lr, + train_step, + wait_for_future, + wait_for_training_service, +) + + +def reset_adapter(train_url, model_id): + """End an existing session so create_model builds fresh factors. + + create_model does NOT reinitialize an existing model_id -- a rerun silently + continues the previous run's adapter, which invalidates any comparison + between runs. Two facts make this awkward: + + * model_id "default" is a RESERVED session and can never be unloaded + ("The default LoRA session is reserved and cannot be unloaded", HTTP 400), + so reruns against it always inherit stale factors. + * every other model_id can be unloaded, and a never-seen id is fresh by + construction. + + Hence each run uses its own id. This only reports what happened; the + frozen-base check on step 1 is what actually enforces freshness. + """ + resp = requests.post( + f"{train_url}/api/v1/unload_model", + json={"model_id": model_id}, + timeout=120, + ) + if resp.status_code in (400, 404, 409): + detail = "" + try: + detail = resp.json().get("detail", "") + except Exception: + pass + return f"not unloaded (HTTP {resp.status_code}: {detail or 'no existing session'})" + resp.raise_for_status() + payload = resp.json() + if "request_id" in payload: + payload = wait_for_future(train_url, payload["request_id"], timeout=600) + return payload + + +def save_adapter(train_url, name): + """Export the trained adapter in the serving layout. + + save_weights_for_sampler writes a LoRA adapter (not a full checkpoint) when + the server runs in an adapter-bearing mode, using the YAML's + lora_export_format — sglang_shared_outer here, the same 3D packed layout the + Qwen reference adapters ship under sglang_shared/. + """ + resp = requests.post( + f"{train_url}/api/v1/save_weights_for_sampler", + json={"model_id": MODEL_ID, "name": name}, # MODEL_ID patched per-run in main() + timeout=120, + ) + resp.raise_for_status() + result = wait_for_future(train_url, resp.json()["request_id"], timeout=1800) + return _raise_on_failed_future(result, "save_weights_for_sampler") + + +def main(): + parser = argparse.ArgumentParser(description="GLM-5.2 password memorization (training only)") + parser.add_argument("--model", type=str, default="zai-org/GLM-5.2-FP8") + parser.add_argument("--train-url", type=str, default="http://localhost:6000") + parser.add_argument("--steps", type=int, default=64) + parser.add_argument("--lr", type=float, default=5e-4) + parser.add_argument( + "--lr-schedule", type=str, default="warmup_cosine", choices=["constant", "cosine", "warmup_cosine"] + ) + parser.add_argument("--lr-min-ratio", type=float, default=0.01) + parser.add_argument("--warmup-steps", type=int, default=8) + parser.add_argument("--log-interval", type=int, default=4) + parser.add_argument( + "--model-id", type=str, default=None, + help=( + "Training session id. Defaults to a per-run id derived from --save-name, because " + "the reserved 'default' session cannot be unloaded and would silently inherit the " + "previous run's adapter." + ), + ) + parser.add_argument( + "--step-timeout", type=float, default=2400.0, + help=( + "Client wait per future. Must exceed the server's own 1800s forward-backward " + "timeout, else the client gives up first and reports a hang the server would " + "have reported itself." + ), + ) + parser.add_argument( + "--repeat", type=int, default=1, + help=( + "Replicate each example N times per step. The 3 password examples pack to ~128 tokens, " + "which is 8 tokens per CP rank at cp_size=16 and leaves nearly all 256 routed experts " + "empty -- the degenerate-batch shape the first run hung on. Raise this to give every " + "CP rank and expert real work." + ), + ) + parser.add_argument("--service-timeout", type=float, default=3600.0, + help="Loading 141 FP8 shards across 16 ranks is slow; default 1h.") + parser.add_argument("--save-name", type=str, default=None, help="Save the adapter under this name when set") + args = parser.parse_args() + + # train_step()/create_model() read MODEL_ID and wait_for_future from their + # defining module at call time, so patching module attributes redirects every + # helper at once. + import run_password_test as _rpt # noqa: PLC0415 + + model_id = args.model_id or (f"run-{args.save_name}" if args.save_name else "run-glm52") + _rpt.MODEL_ID = model_id + globals()["MODEL_ID"] = model_id + + _orig_wait = _rpt.wait_for_future + _rpt.wait_for_future = lambda url, rid, timeout=args.step_timeout: _orig_wait(url, rid, timeout=timeout) + + print(f" Waiting for the training server at {args.train_url} (timeout {args.service_timeout:.0f}s)...") + if not wait_for_training_service(args.train_url, timeout=args.service_timeout): + print(" FAILED: training server never reported engine_running") + return 1 + print(" Training server ready.") + + tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True) + training_data = build_training_data(tokenizer) * args.repeat + tokens = sum(len(d["model_input"]["input_ids"]) for d in training_data) + print(f" Built {len(training_data)} examples over {len(CODES)} project codes " + f"(repeat={args.repeat}, {tokens} tokens total, ~{tokens // 16} per CP rank).") + + print(f" Session id: {model_id}") + print(f" Reset: {reset_adapter(args.train_url, model_id)}") + create_result = create_model(args.train_url, args.model) + print(f" Model created: model_id={create_result.get('model_id', MODEL_ID)}") + + print(f"\n Training ({args.steps} steps, lr={args.lr}, schedule={args.lr_schedule})...") + t0 = time.time() + first_loss = last_loss = None + for step in range(args.steps): + step_lr = get_lr(step, args) + loss, grad_norm = train_step(args.train_url, training_data, step_lr) + if first_loss is None: + first_loss = loss + # A fresh adapter has lora_B == 0, so step 1 MUST equal the frozen + # base loss. Anything else means we resumed stale factors. + if isinstance(loss, (int, float)) and loss < 1.0: + print( + f"\n ERROR: step 1 loss {loss:.4f} is far below the frozen-base value; " + "the adapter was NOT reset and this run continues stale factors." + ) + return 1 + last_loss = loss + step_num = step + 1 + if step_num == 1 or step_num == args.steps or step_num % args.log_interval == 0: + print(f" Step {step_num}/{args.steps}: loss={loss}, grad_norm={grad_norm}, lr={step_lr:.2e}") + print(f" Training done in {time.time() - t0:.1f}s (loss {first_loss} -> {last_loss})") + + if args.save_name: + result = save_adapter(args.train_url, args.save_name) + print(f" Adapter saved: {result}") + + # The Qwen reference adapters reach ~0 loss within 16 steps on this task. + if isinstance(last_loss, (int, float)) and last_loss > 0.1: + print(f"\n WARNING: final loss {last_loss} is above 0.1 — memorization did not converge.") + return 1 + print("\n Training-half PASSED. Recall verification needs a GLM-5.2 SGLang endpoint (see the runbook).") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/src/xorl/lora/utils.py b/src/xorl/lora/utils.py index 3307cbcb..8a73fa16 100644 --- a/src/xorl/lora/utils.py +++ b/src/xorl/lora/utils.py @@ -232,10 +232,18 @@ def _find_target_modules( Uses the LoRA mapping registry to determine which modules can have LoRA applied to them. Supports two matching modes: - 1. Direct match: Module name matches a target (e.g., "q_proj", "experts") - 2. Indirect match: Module has children matching targets (e.g., MoE experts + 1. Path match: a target containing a path separator or glob metacharacter + (e.g. "*.mlp.experts.*.gate_proj") is matched with fnmatch against the + FULL module path. Leaf matching discards a module's position in the tree, + so a bare "gate_proj" cannot distinguish a routed expert from a shared + expert or a dense MLP; a path pattern can. + 2. Direct match: Module name matches a target (e.g., "q_proj", "experts") + 3. Indirect match: Module has children matching targets (e.g., MoE experts module has "gate_proj", "up_proj", "down_proj" as weight attributes) + Path patterns participate in the same coverage check as bare names: a + pattern matching nothing raises rather than silently adapting less. + The algorithm processes modules top-down and skips children of replaced modules to avoid double-replacement. @@ -256,6 +264,11 @@ def _find_target_modules( replaced_prefixes: Set[str] = set() # Track replaced module paths to skip their children matched_targets: Set[str] = set(satisfied_targets or ()) + # A target is a path pattern when it carries positional information. Bare + # names keep the historical leaf/indirect semantics untouched. + path_patterns = [t for t in target_modules if any(ch in t for ch in ".*?[")] + leaf_targets = [t for t in target_modules if t not in path_patterns] + for name, module in model.named_modules(): # Skip if this module is under an already-matched parent # (avoid replacing children of modules we're going to replace) @@ -268,8 +281,16 @@ def _find_target_modules( module_name = name.split(".")[-1] if name else "" + # Path match: explicit about position, so it wins over leaf matching. + matched_pattern = next((p for p in path_patterns if fnmatch.fnmatch(name, p)), None) + if matched_pattern is not None: + matched_paths.append(name) + replaced_prefixes.add(name) + matched_targets.add(matched_pattern) + continue + # Direct match: module name matches target_modules - if module_name in target_modules: + if module_name in leaf_targets: matched_paths.append(name) replaced_prefixes.add(name) matched_targets.add(module_name) @@ -279,7 +300,7 @@ def _find_target_modules( # This handles MoE experts where user specifies "gate_proj" but the # actual module to replace is "experts" which contains gate_proj weights module_attrs = set(dir(module)) - indirect_matches = {target for target in target_modules if target in module_attrs} + indirect_matches = {target for target in leaf_targets if target in module_attrs} if indirect_matches: matched_paths.append(name) replaced_prefixes.add(name) diff --git a/src/xorl/models/auto.py b/src/xorl/models/auto.py index 1085225a..cf04f9f7 100644 --- a/src/xorl/models/auto.py +++ b/src/xorl/models/auto.py @@ -817,6 +817,7 @@ def build_foundation_model( server_training: bool = False, enable_lora: bool = False, block_fp8_qlora_training: bool = False, + glm52_lora_scope: str = "all", glm52_fullparam_fp8_training: bool = False, lora_rank: Optional[int] = None, lora_alpha: Optional[int] = None, @@ -853,13 +854,22 @@ def build_foundation_model( raise ValueError( "glm52_fullparam_fp8_training and block_fp8_qlora_training are mutually exclusive training lanes" ) - exact_active_lora = bool(server_training and glm52_model and block_fp8_qlora_training and ep_dispatch == "alltoall") + # Scope selects WHICH FACTORS TRAIN, not which modules are adapted: the + # complete exact family is always constructed so gradients can flow through + # every region (NativeBlockFP8Linear is forward-only). Out-of-scope factors + # are frozen after construction and keep lora_B == 0, so the forward program + # is identical for every scope. + glm52_lora_scope = str(glm52_lora_scope or "all") + exact_active_lora = bool( + server_training and glm52_model and block_fp8_qlora_training and ep_dispatch == "alltoall" + ) if exact_active_lora: glm52_exact_lora_scaling(lora_rank, lora_alpha) # Training lanes select the same exact value family through their own # admission flags. The scoring-only flag must remain off for either one: # it describes a frozen trunk, which neither training admission permits. config._glm52_block_fp8_qlora = bool(block_fp8_qlora_training) + config._glm52_lora_scope = glm52_lora_scope config._glm52_fullparam_training = bool(glm52_fullparam_fp8_training) config._glm52_exact_contract = bool( server_training and glm52_model and not block_fp8_qlora_training and not glm52_fullparam_fp8_training diff --git a/src/xorl/models/transformers/glm5/exact_lm_head_qlora.py b/src/xorl/models/transformers/glm5/exact_lm_head_qlora.py index 2eb74b5c..85827a57 100644 --- a/src/xorl/models/transformers/glm5/exact_lm_head_qlora.py +++ b/src/xorl/models/transformers/glm5/exact_lm_head_qlora.py @@ -828,8 +828,16 @@ def _validate_operands( ) if local_weight.requires_grad: raise RuntimeError("GLM-5.2 exact LM-head base weight must remain frozen") - if not lora_A.requires_grad or not local_lora_B.requires_grad: - raise RuntimeError("GLM-5.2 exact LM-head A and B factor masters must both be trainable") + if lora_A.requires_grad != local_lora_B.requires_grad: + # The invariant is that the pair SHARES trainability -- a half-frozen + # pair would take a gradient on one side of a product whose other + # side is fixed. Both-frozen is consistent and is what a narrowed + # glm52_lora_scope produces: the lm-head module is still constructed + # so gradients flow THROUGH it, but its factors do not learn. + raise RuntimeError( + "GLM-5.2 exact LM-head A and B factor masters must share trainability, got " + f"A.requires_grad={lora_A.requires_grad}, B.requires_grad={local_lora_B.requires_grad}" + ) expected_strides = { "local_weight": (GLM52_LM_HEAD_HIDDEN_SIZE, 1), diff --git a/src/xorl/models/transformers/glm5/qlora.py b/src/xorl/models/transformers/glm5/qlora.py index 0786392a..5b09d2f4 100644 --- a/src/xorl/models/transformers/glm5/qlora.py +++ b/src/xorl/models/transformers/glm5/qlora.py @@ -2,6 +2,7 @@ from __future__ import annotations +import logging from collections import Counter from dataclasses import dataclass @@ -38,6 +39,9 @@ from xorl.qlora.modules.moe_experts import BlockFP8QLoRAMoeExperts +logger = logging.getLogger(__name__) + + GLM52_QLORA_ORDINARY_TARGET_COUNT = 625 GLM52_QLORA_QUANTIZED_LINEAR_COUNT = 624 GLM52_QLORA_ROUTED_BANK_COUNT = 75 @@ -52,6 +56,52 @@ "o_proj", ) _MLP_TARGETS = ("gate_proj", "up_proj", "down_proj") + +# Adapted-target scopes. "all" is the complete deterministic inventory and the +# ONLY scope docs/k3/LORA_CONTRACT.md qualifies for train/serve bit-exactness. +# The narrower scopes exist to isolate where MoE learning happens; they select a +# different target universe, so the exact active-LoRA family is refused for them +# (see auto.py) and they run the ordinary block-FP8 lane. Modules left out of the +# target list are converted to frozen native FP8 by replace_glm52_native_fp8_modules. +GLM52_LORA_SCOPES = ("all", "moe", "shared_experts", "routed_experts") + + +def glm52_scoped_factor_names(inventory: "Glm52AdapterInventory", scope: str) -> set[str]: + """Factor names a scope leaves trainable.""" + return {factor.name for factor in inventory.factors if glm52_scope_admits(scope, factor.role)} + + +def apply_glm52_lora_scope(model: nn.Module, inventory: "Glm52AdapterInventory", scope: str) -> int: + """Freeze every factor outside ``scope``; return how many were frozen. + + The complete adapter inventory is always constructed: NativeBlockFP8Linear is + forward-only ("phase-one forward is scoring-only; activation backward + requires a validated kernel"), so any region left unadapted would block + gradients from reaching adapted regions downstream of it. Keeping every exact + module and freezing the factors instead gives a working backward everywhere, + while a frozen pair keeps lora_B == 0 and so contributes nothing to the + forward. + """ + trainable = glm52_scoped_factor_names(inventory, scope) + frozen = 0 + for name, parameter in model.named_parameters(): + if name in inventory.factor_names and name not in trainable: + parameter.requires_grad = False + frozen += 1 + return frozen + + +def glm52_scope_admits(scope: str, role: str) -> bool: + """Return whether a target role belongs to the selected scope.""" + if scope == "all": + return True + if scope == "moe": + return role.startswith("shared_expert.") or role == "routed_expert" + if scope == "shared_experts": + return role.startswith("shared_expert.") + if scope == "routed_experts": + return role == "routed_expert" + raise ValueError(f"Unknown GLM-5.2 LoRA scope: {scope!r} (expected one of {GLM52_LORA_SCOPES})") _EXPERT_FACTORS = ( "gate_proj_lora_A", "gate_proj_lora_B", @@ -154,20 +204,25 @@ def _validate_official_config(config) -> dict: exact_shared_component = bool(getattr(config, "_glm52_exact_active_lora_shared_expert_component", False)) exact_routed_component = bool(getattr(config, "_glm52_exact_active_lora_routed_expert_component", False)) exact_lm_head_component = bool(getattr(config, "_glm52_exact_active_lora_lm_head_component", False)) - if exact_attention_component and not exact_dense_component: - raise ValueError("GLM-5.2 exact active-LoRA attention component requires the exact active-LoRA dense component") - if exact_shared_component and not exact_attention_component: - raise ValueError( - "GLM-5.2 exact active-LoRA shared-expert component requires the exact active-LoRA attention component" - ) - if exact_routed_component and not exact_shared_component: - raise ValueError( - "GLM-5.2 exact active-LoRA routed-expert component requires the exact active-LoRA shared-expert component" - ) - if exact_lm_head_component and not exact_routed_component: - raise ValueError( - "GLM-5.2 exact active-LoRA lm-head component requires the exact active-LoRA routed-expert component" - ) + # The complete family is one indivisible value program, so its components are + # strictly nested: no member may be enabled without its predecessors. + if True: + if exact_attention_component and not exact_dense_component: + raise ValueError( + "GLM-5.2 exact active-LoRA attention component requires the exact active-LoRA dense component" + ) + if exact_shared_component and not exact_attention_component: + raise ValueError( + "GLM-5.2 exact active-LoRA shared-expert component requires the exact active-LoRA attention component" + ) + if exact_routed_component and not exact_shared_component: + raise ValueError( + "GLM-5.2 exact active-LoRA routed-expert component requires the exact active-LoRA shared-expert component" + ) + if exact_lm_head_component and not exact_routed_component: + raise ValueError( + "GLM-5.2 exact active-LoRA lm-head component requires the exact active-LoRA routed-expert component" + ) exact_component_enabled = any( ( exact_dense_component, @@ -294,6 +349,7 @@ def _expected_targets(model: nn.Module, config) -> tuple[Glm52AdapterTarget, ... if lm_head.weight.dtype is not torch.bfloat16: raise TypeError(f"GLM-5.2 QLoRA lm_head base must remain BF16, got {lm_head.weight.dtype}") targets.append(Glm52AdapterTarget("lm_head", "output.lm_head", "bf16_linear", *lm_head_shape)) + return tuple(targets) @@ -575,6 +631,9 @@ def _build_factor_inventory( def _validate_constructed_model(model: nn.Module, inventory: Glm52AdapterInventory) -> None: + # Bound first: both the native-FP8 check and the geometry assertions below + # branch on it. + scope = str(getattr(model.config, "_glm52_lora_scope", "all") or "all") expected_quantized = {target.name for target in inventory.targets if target.kind == "block_fp8_linear"} expected_heads = {target.name for target in inventory.targets if target.kind == "bf16_linear"} expected_banks = {target.name for target in inventory.targets if target.kind == "block_fp8_routed_bank"} @@ -685,6 +744,8 @@ def _validate_constructed_model(model: nn.Module, inventory: Glm52AdapterInvento and not isinstance(module, (Glm52ExactTP1DenseMLP, Glm52ExactTP1AbsorbedKvBBlockFP8QLoRA)) and not any(name.startswith(f"{root}.") for root in expected_exact_shared_roots) } + # Every non-indexer module is adapted in every scope, so the native-FP8 set is + # exactly the DSA selector projections. if actual_native_indexers != expected_native_indexers: raise RuntimeError( "GLM-5.2 QLoRA frozen native indexer set mismatch: " @@ -696,20 +757,16 @@ def _validate_constructed_model(model: nn.Module, inventory: Glm52AdapterInvento "GLM-5.2 QLoRA requires exactly " f"{GLM52_FROZEN_NATIVE_INDEXER_PROJECTION_COUNT} frozen native indexer projections" ) - for layer_idx, indexer_type in enumerate(_official_indexer_schedule()): - if indexer_type != "full": - continue - weights_proj = model.get_submodule(f"model.layers.{layer_idx}.self_attn.indexer.weights_proj") - if type(weights_proj) is not nn.Linear or weights_proj.weight.dtype is not torch.bfloat16: - raise RuntimeError( - f"GLM-5.2 QLoRA indexer weights_proj for layer {layer_idx} must remain an ordinary BF16 linear" - ) + # Scope selects WHICH FACTORS TRAIN, not which modules exist: every region + # keeps its exact adapter module so gradients can flow through it. Frozen + # factors keep lora_B == 0, so they contribute nothing to the forward. + expected_trainable = glm52_scoped_factor_names(inventory, scope) trainable = {name for name, parameter in model.named_parameters() if parameter.requires_grad} - if trainable != inventory.factor_names: + if trainable != expected_trainable: raise RuntimeError( - "GLM-5.2 QLoRA trainable factor set mismatch: " - f"missing={sorted(inventory.factor_names - trainable)} extra={sorted(trainable - inventory.factor_names)}" + f"GLM-5.2 QLoRA trainable factor set mismatch under scope {scope!r}: " + f"missing={sorted(expected_trainable - trainable)} extra={sorted(trainable - expected_trainable)}" ) trainable_parameters = [parameter for parameter in model.parameters() if parameter.requires_grad] if len({id(parameter) for parameter in trainable_parameters}) != len(trainable_parameters): @@ -871,12 +928,25 @@ def prepare_glm52_block_fp8_qlora( replace_glm52_native_fp8_modules(model, quantization_config) inventory = Glm52AdapterInventory(targets=targets, factors=_build_factor_inventory(model, targets)) + scope = str(getattr(config, "_glm52_lora_scope", "all") or "all") + if scope not in GLM52_LORA_SCOPES: + raise ValueError(f"Unknown GLM-5.2 LoRA scope: {scope!r} (expected one of {GLM52_LORA_SCOPES})") + frozen = apply_glm52_lora_scope(model, inventory, scope) + if scope != "all": + logger.info( + f"GLM-5.2 LoRA scope {scope!r}: froze {frozen} of {len(inventory.factors)} factors; " + f"{len(inventory.factors) - frozen} remain trainable" + ) _validate_constructed_model(model, inventory) model._glm52_adapter_inventory = inventory return inventory __all__ = [ + "GLM52_LORA_SCOPES", + "apply_glm52_lora_scope", + "glm52_scope_admits", + "glm52_scoped_factor_names", "GLM52_FROZEN_NATIVE_INDEXER_PROJECTION_COUNT", "GLM52_QLORA_FACTOR_COUNT", "GLM52_QLORA_ORDINARY_TARGET_COUNT", diff --git a/src/xorl/server/runner/adapters/manager.py b/src/xorl/server/runner/adapters/manager.py index b13d640a..e67b5fa7 100644 --- a/src/xorl/server/runner/adapters/manager.py +++ b/src/xorl/server/runner/adapters/manager.py @@ -2593,7 +2593,12 @@ def stage_gradient_numerators( for item in plan.parameters: parameter = named_parameters[item.fqn] if parameter.grad is None: - if item.requires_local_gradient: + # A frozen factor legitimately produces no gradient. glm52_lora_scope + # builds the complete adapter inventory -- every region needs a module + # with a working backward -- and then freezes the factors outside the + # scope, so "absent" is only a fault for a factor that is supposed to + # be learning. + if item.requires_local_gradient and parameter.requires_grad: raise AdapterGradientOwnershipError(f"Required adapter gradient is absent for {item.fqn!r}") continue local_gradient = self._capture_local_gradient(state, item, parameter) @@ -2637,8 +2642,14 @@ def commit_gradient_capture(self, model_id: str) -> tuple[int, int]: raise AdapterGradientOwnershipError("Staged gradient denominator must be finite and positive") if scratch.staged_numerator_scale is None or not math.isfinite(scratch.staged_numerator_scale): raise AdapterGradientOwnershipError("Staged gradient numerator scale is missing or nonfinite") + frozen_fqns = { + canonical_parameter_name(name) + for name, parameter in self.model.named_parameters() + if not parameter.requires_grad + } for fqn, item in item_by_fqn.items(): - if item.requires_local_gradient and fqn not in staged_set: + # Frozen factors (glm52_lora_scope) stage nothing by design. + if item.requires_local_gradient and fqn not in staged_set and fqn not in frozen_fqns: raise AdapterGradientOwnershipError(f"Required staged adapter gradient is absent for {fqn!r}") if fqn not in staged_set: continue @@ -2835,10 +2846,16 @@ def _validate_authoritative_state(self, state: AdapterState, *, allow_optimizer_ raise AdapterGradientOwnershipError("Compiled producer and captured gradient source disagree") scratch_by_fqn = {canonical_parameter_name(name): tensor for name, tensor in scratch.numerators.items()} local_by_fqn = {canonical_parameter_name(name): parameter for name, parameter in state.local_params.items()} + frozen_fqns = { + canonical_parameter_name(name) + for name, parameter in self.model.named_parameters() + if not parameter.requires_grad + } for item in plan.parameters: numerator = scratch_by_fqn.get(item.fqn) if numerator is None: - if item.requires_local_gradient: + # Frozen factors (glm52_lora_scope) produce no numerator by design. + if item.requires_local_gradient and item.fqn not in frozen_fqns: raise AdapterGradientOwnershipError(f"Required numerator is absent for {item.fqn!r}") continue if tuple(numerator.shape) != tuple(local_by_fqn[item.fqn].shape): diff --git a/src/xorl/server/runner/model_runner.py b/src/xorl/server/runner/model_runner.py index bb17cac6..d27639e8 100644 --- a/src/xorl/server/runner/model_runner.py +++ b/src/xorl/server/runner/model_runner.py @@ -1593,6 +1593,7 @@ def _initialize_model(self): moe_hybrid_shared_lora=self.lora_config.get("moe_hybrid_shared_lora", False), enable_qlora=enable_qlora, block_fp8_qlora_training=block_fp8_qlora_training, + glm52_lora_scope=self.lora_config.get("glm52_lora_scope", "all"), quant_format=self.lora_config.get("quant_format", "nvfp4"), quant_group_size=self.lora_config.get("quant_group_size", 16), qlora_exclude_modules=self.lora_config.get("exclude_modules"), diff --git a/src/xorl/server/runner/setup.py b/src/xorl/server/runner/setup.py index 629f2c95..ba6e7863 100644 --- a/src/xorl/server/runner/setup.py +++ b/src/xorl/server/runner/setup.py @@ -16,6 +16,7 @@ import torch import torch.distributed as dist import yaml +from torch.distributed.elastic.multiprocessing.errors import record as _elastic_record from xorl.arguments import Arguments, parse_args from xorl.server.runner.model_runner import ModelRunner @@ -146,6 +147,14 @@ async def _run_worker(config: Dict[str, Any], bind_address: str, output_dir: str log_level: Logging level string (e.g. "INFO", "DEBUG") """ # Set NCCL environment variables BEFORE initializing any NCCL operations + # Native-crash diagnostics: workers were exiting with code 70 and no Python + # traceback, so the fault is below Python. faulthandler dumps a stack on + # fatal signals; TORCH_SHOW_CPP_STACKTRACES surfaces the C++ frames behind a + # torch error instead of a bare message. + import faulthandler # noqa: PLC0415 + + faulthandler.enable(all_threads=True) + os.environ.setdefault("TORCH_SHOW_CPP_STACKTRACES", "1") os.environ.setdefault("NCCL_CUMEM_ENABLE", "0") os.environ.setdefault("NCCL_NVLS_ENABLE", "0") os.environ.setdefault("TORCH_NCCL_BLOCKING_WAIT", "1") @@ -267,10 +276,25 @@ def _is_server_args_config(config_path: str) -> bool: return bool(server_args_keys & set(config.keys())) +@_elastic_record def main(): - """Main entry point with command-line argument parsing.""" + """Main entry point with command-line argument parsing. + + Decorated with torch.distributed.elastic's ``record`` so a worker failure + writes its Python traceback to TORCHELASTIC_ERROR_FILE. Without it torchrun + reports only ``error_file: `` and an exit code, which is not enough to + diagnose a rank that dies inside the model forward. + """ # Set NCCL environment variables BEFORE any NCCL initialization # These must be set before torch.distributed is imported/initialized + # Native-crash diagnostics: workers were exiting with code 70 and no Python + # traceback, so the fault is below Python. faulthandler dumps a stack on + # fatal signals; TORCH_SHOW_CPP_STACKTRACES surfaces the C++ frames behind a + # torch error instead of a bare message. + import faulthandler # noqa: PLC0415 + + faulthandler.enable(all_threads=True) + os.environ.setdefault("TORCH_SHOW_CPP_STACKTRACES", "1") os.environ.setdefault("NCCL_CUMEM_ENABLE", "0") os.environ.setdefault("NCCL_NVLS_ENABLE", "0") os.environ.setdefault("TORCH_NCCL_BLOCKING_WAIT", "1") diff --git a/src/xorl/server/server_arguments.py b/src/xorl/server/server_arguments.py index 4af38c9a..e72ce34b 100644 --- a/src/xorl/server/server_arguments.py +++ b/src/xorl/server/server_arguments.py @@ -18,6 +18,9 @@ from xorl.ops.loss import CrossEntropyMode +GLM52_LORA_SCOPE_CHOICES = ("all", "moe", "shared_experts", "routed_experts") + + @dataclass class ServerArguments: """ @@ -1129,6 +1132,20 @@ class ServerArguments: default=False, metadata={"help": "Enable QLoRA (quantized LoRA) for memory-efficient training"} ) + glm52_lora_scope: str = field( + default="all", + metadata={ + "help": ( + "Which GLM-5.2 modules receive LoRA factors. 'all' (default) is the complete " + "deterministic inventory and the only scope qualified for train/serve bit-exactness. " + "'moe' trains shared + routed experts, 'shared_experts' and 'routed_experts' isolate one " + "each. Scope selects which factors TRAIN, not which modules are adapted: the complete " + "exact family is always built so gradients can flow through every region, and " + "out-of-scope factors are frozen with lora_B == 0 so they never affect the forward." + ) + }, + ) + block_fp8_qlora_training: bool = field( default=False, metadata={ @@ -1376,8 +1393,22 @@ def __post_init__(self): ] if mismatches: raise ValueError("GLM-5.2 block-FP8 QLoRA rejects unsupported configuration: " + ", ".join(mismatches)) + if self.glm52_lora_scope not in GLM52_LORA_SCOPE_CHOICES: + raise ValueError( + f"glm52_lora_scope must be one of {GLM52_LORA_SCOPE_CHOICES}, got {self.glm52_lora_scope!r}" + ) + if self.glm52_lora_scope != "all" and self.ep_dispatch != "alltoall": + # Scope only narrows WHICH FACTORS TRAIN. The complete exact + # family is still constructed, so the exact transport applies. + raise ValueError( + f"glm52_lora_scope={self.glm52_lora_scope!r} still builds the complete exact family and " + f"requires ep_dispatch='alltoall', got {self.ep_dispatch!r}" + ) if self.lora_target_modules is not None or self.lora_target_manifest is not None: - raise ValueError("GLM-5.2 block-FP8 QLoRA uses its complete deterministic target set") + raise ValueError( + "GLM-5.2 block-FP8 QLoRA uses its deterministic target set; " + "select a subset with glm52_lora_scope instead" + ) if self.qlora_exclude_modules is not None: raise ValueError("GLM-5.2 block-FP8 QLoRA derives checkpoint exclusions and rejects user overrides") if self.merge_lora_interval: @@ -1643,6 +1674,7 @@ def to_config_dict(self) -> Dict[str, Any]: "lora_export_format": self.lora_export_format, "enable_qlora": self.enable_qlora, "block_fp8_qlora_training": self.block_fp8_qlora_training, + "glm52_lora_scope": self.glm52_lora_scope, "quant_format": self.quant_format, "quant_group_size": self.quant_group_size, "exclude_modules": self.qlora_exclude_modules, diff --git a/src/xorl/trainers/model_builder.py b/src/xorl/trainers/model_builder.py index cf224676..44d90729 100644 --- a/src/xorl/trainers/model_builder.py +++ b/src/xorl/trainers/model_builder.py @@ -221,6 +221,7 @@ def build_training_model( # --- QLoRA --- enable_qlora: bool = False, block_fp8_qlora_training: bool = False, + glm52_lora_scope: str = "all", quant_format: str = "nvfp4", quant_group_size: int = 16, qlora_exclude_modules: Optional[List[str]] = None, @@ -406,6 +407,7 @@ def build_training_model( server_training=server_training, enable_lora=enable_lora, block_fp8_qlora_training=block_fp8_qlora_training, + glm52_lora_scope=glm52_lora_scope, glm52_fullparam_fp8_training=glm52_fullparam_fp8_training, lora_rank=lora_rank, lora_alpha=lora_alpha, diff --git a/tests/models/test_glm52_lora_scope.py b/tests/models/test_glm52_lora_scope.py new file mode 100644 index 00000000..f3f41ad1 --- /dev/null +++ b/tests/models/test_glm52_lora_scope.py @@ -0,0 +1,142 @@ +"""Scoped GLM-5.2 LoRA target selection. + +The default scope reproduces the complete deterministic inventory; the narrowed +scopes isolate shared and/or routed experts so the contribution of each can be +measured. Narrowed scopes select a different target universe than the qualified +exact active-LoRA family, so they must refuse it. +""" + +import pytest +import torch + +from xorl.models.transformers.glm5.qlora import GLM52_LORA_SCOPES, glm52_scope_admits + + +pytestmark = [pytest.mark.cpu] + + +# One representative role per region of the official inventory. +_ROLES = ( + "attention.q_a_proj", + "attention.o_proj", + "dense_mlp.gate_proj", + "shared_expert.gate_proj", + "shared_expert.down_proj", + "routed_expert", + "output.lm_head", +) + + +def test_scope_names_are_stable(): + assert GLM52_LORA_SCOPES == ("all", "moe", "shared_experts", "routed_experts") + + +def test_all_scope_admits_every_region(): + assert all(glm52_scope_admits("all", role) for role in _ROLES) + + +@pytest.mark.parametrize( + "scope,expected", + [ + ("moe", {"shared_expert.gate_proj", "shared_expert.down_proj", "routed_expert"}), + ("shared_experts", {"shared_expert.gate_proj", "shared_expert.down_proj"}), + ("routed_experts", {"routed_expert"}), + ], +) +def test_narrowed_scopes_select_only_their_region(scope, expected): + admitted = {role for role in _ROLES if glm52_scope_admits(scope, role)} + assert admitted == expected + + +@pytest.mark.parametrize("scope", ["moe", "shared_experts", "routed_experts"]) +def test_narrowed_scopes_never_admit_attention_or_head(scope): + """Attention and the lm_head are what distinguish these from the full set.""" + assert not glm52_scope_admits(scope, "attention.q_a_proj") + assert not glm52_scope_admits(scope, "attention.kv_b_proj") + assert not glm52_scope_admits(scope, "output.lm_head") + assert not glm52_scope_admits(scope, "dense_mlp.gate_proj") + + +def test_unknown_scope_fails_closed(): + with pytest.raises(ValueError, match="Unknown GLM-5.2 LoRA scope"): + glm52_scope_admits("experts_only", "routed_expert") + + +def test_server_arguments_reject_unknown_scope(): + from xorl.server.server_arguments import GLM52_LORA_SCOPE_CHOICES + + assert GLM52_LORA_SCOPE_CHOICES == GLM52_LORA_SCOPES + + +# --------------------------------------------------------------------------- +# Scope selects trainability, not construction +# --------------------------------------------------------------------------- + + +def test_scoped_factor_names_partition_the_inventory(): + """Every factor is either trainable under a scope or frozen -- never dropped. + + The complete inventory is always built: NativeBlockFP8Linear is forward-only, + so a region left unadapted would block gradients from reaching adapted + regions downstream of it. + """ + from types import SimpleNamespace + + from xorl.models.transformers.glm5.qlora import glm52_scoped_factor_names + + factors = [ + SimpleNamespace(name=f"f{i}", role=role) + for i, role in enumerate( + ["attention.q_a_proj", "dense_mlp.gate_proj", "shared_expert.up_proj", "routed_expert", "output.lm_head"] + ) + ] + inventory = SimpleNamespace(factors=factors) + + assert glm52_scoped_factor_names(inventory, "all") == {"f0", "f1", "f2", "f3", "f4"} + assert glm52_scoped_factor_names(inventory, "moe") == {"f2", "f3"} + assert glm52_scoped_factor_names(inventory, "shared_experts") == {"f2"} + assert glm52_scoped_factor_names(inventory, "routed_experts") == {"f3"} + + +def test_freezing_is_the_complement_of_the_scope(): + """apply_glm52_lora_scope must freeze exactly the out-of-scope factors.""" + from types import SimpleNamespace + + import torch.nn as nn + + from xorl.models.transformers.glm5.qlora import apply_glm52_lora_scope + + model = nn.Module() + for name in ("attn_f", "shared_f", "routed_f"): + setattr(model, name, nn.Parameter(torch.zeros(2))) + inventory = SimpleNamespace( + factors=[ + SimpleNamespace(name="attn_f", role="attention.q_a_proj"), + SimpleNamespace(name="shared_f", role="shared_expert.up_proj"), + SimpleNamespace(name="routed_f", role="routed_expert"), + ], + factor_names=frozenset({"attn_f", "shared_f", "routed_f"}), + ) + + frozen = apply_glm52_lora_scope(model, inventory, "routed_experts") + assert frozen == 2 + assert model.attn_f.requires_grad is False + assert model.shared_f.requires_grad is False + assert model.routed_f.requires_grad is True + + +def test_scope_all_freezes_nothing(): + from types import SimpleNamespace + + import torch.nn as nn + + from xorl.models.transformers.glm5.qlora import apply_glm52_lora_scope + + model = nn.Module() + model.routed_f = nn.Parameter(torch.zeros(2)) + inventory = SimpleNamespace( + factors=[SimpleNamespace(name="routed_f", role="routed_expert")], + factor_names=frozenset({"routed_f"}), + ) + assert apply_glm52_lora_scope(model, inventory, "all") == 0 + assert model.routed_f.requires_grad is True diff --git a/tests/models/test_lora_path_pattern_targets.py b/tests/models/test_lora_path_pattern_targets.py new file mode 100644 index 00000000..2a4bf5c2 --- /dev/null +++ b/tests/models/test_lora_path_pattern_targets.py @@ -0,0 +1,108 @@ +"""Path-qualified LoRA target selection. + +Leaf-name matching discards a module's position in the tree, so `gate_proj` +selects routed experts, shared experts, and dense MLPs alike -- there is no way +to adapt one without the others. Targets containing a path separator or glob +are matched against the FULL module path instead, which makes those regions +separable on any architecture. + +Bare names must keep behaving exactly as before. +""" + +import pytest +import torch.nn as nn + +from xorl.lora.utils import _find_target_modules + + +pytestmark = [pytest.mark.cpu] + + +def _proj_block(hidden=8, inter=16): + block = nn.Module() + block.gate_proj = nn.Linear(hidden, inter, bias=False) + block.up_proj = nn.Linear(hidden, inter, bias=False) + block.down_proj = nn.Linear(inter, hidden, bias=False) + return block + + +def _moe_model(num_layers=2, num_experts=3): + """Tree mirroring a real MoE layout: routed + shared + attention.""" + model = nn.Module() + model.layers = nn.ModuleList() + for _ in range(num_layers): + layer = nn.Module() + layer.self_attn = nn.Module() + layer.self_attn.q_proj = nn.Linear(8, 8, bias=False) + layer.self_attn.o_proj = nn.Linear(8, 8, bias=False) + layer.mlp = nn.Module() + layer.mlp.experts = nn.ModuleList([_proj_block() for _ in range(num_experts)]) + layer.mlp.shared_expert = _proj_block() + model.layers.append(layer) + return model + + +def test_bare_names_are_unchanged_and_hit_every_region(): + """Regression: the pre-existing behaviour must be byte-identical.""" + model = _moe_model() + paths = _find_target_modules(model, ["gate_proj"]) + assert any(".mlp.experts." in p for p in paths), "routed experts not covered" + assert any(".mlp.shared_expert." in p for p in paths), "shared expert not covered" + # 2 layers x (3 routed + 1 shared) + assert len(paths) == 8 + + +def test_path_pattern_selects_routed_experts_only(): + model = _moe_model() + paths = _find_target_modules(model, ["*.mlp.experts.*.gate_proj"]) + assert len(paths) == 6 # 2 layers x 3 experts + assert all(".mlp.experts." in p for p in paths) + assert not any("shared_expert" in p for p in paths) + + +def test_path_pattern_selects_shared_expert_only(): + model = _moe_model() + paths = _find_target_modules(model, ["*.shared_expert.*_proj"]) + assert len(paths) == 6 # 2 layers x 3 projections + assert all("shared_expert" in p for p in paths) + assert not any(".mlp.experts." in p for p in paths) + + +def test_patterns_and_bare_names_compose(): + model = _moe_model() + paths = _find_target_modules(model, ["q_proj", "*.shared_expert.down_proj"]) + assert sum("q_proj" in p for p in paths) == 2 + assert sum("shared_expert.down_proj" in p for p in paths) == 2 + assert not any(".mlp.experts." in p for p in paths) + + +def test_unmatched_pattern_fails_closed(): + """A pattern that selects nothing must raise, like an unmatched bare name.""" + model = _moe_model() + with pytest.raises(ValueError, match="matched no module"): + _find_target_modules(model, ["*.mlp.nonexistent.*"]) + + +def test_unmatched_bare_name_still_fails_closed(): + model = _moe_model() + with pytest.raises(ValueError, match="matched no module"): + _find_target_modules(model, ["v_proj"]) + + +def test_matched_paths_never_nest(): + """No matched path may be an ancestor of another, else it is replaced twice. + + fnmatch's ``*`` spans ``.``, so "*.mlp.experts.*" matches both an expert + container and the projections beneath it. Only LoRA-applicable modules are + considered, and a replaced parent suppresses its children, so the result + must still be a disjoint set. + """ + model = _moe_model() + paths = _find_target_modules(model, ["*.mlp.experts.*"]) + assert paths, "pattern selected nothing" + assert all( + not other.startswith(path + ".") for path in paths for other in paths if other != path + ), f"nested matches would be double-replaced: {paths}" + # 2 layers x 3 experts x 3 projections; the bare-nn.Module expert blocks are + # not LoRA-applicable, so the applicable descendants are selected instead. + assert len(paths) == 18 From dacf32621d257e3abc0a476003bc133d12c331c8 Mon Sep 17 00:00:00 2001 From: Qingyang Wu Date: Wed, 19 Aug 2026 22:51:30 +0000 Subject: [PATCH 2/2] Support single-pair training and result.json in the GLM-5.2 password driver The published password-adapter repos train one adapter per password, so add --project/--password to narrow the shared CODES table to a single pair, and --result-json to emit {project, password, final_loss, train_time_sec} beside the exported adapter in the same shape those repos ship. --- .../run_glm52_password_train.py | 34 ++++++++++++++++++- 1 file changed, 33 insertions(+), 1 deletion(-) diff --git a/examples/server/password_memorization/run_glm52_password_train.py b/examples/server/password_memorization/run_glm52_password_train.py index 0d0d1b53..54c8afae 100644 --- a/examples/server/password_memorization/run_glm52_password_train.py +++ b/examples/server/password_memorization/run_glm52_password_train.py @@ -114,6 +114,18 @@ def main(): parser.add_argument("--lr-min-ratio", type=float, default=0.01) parser.add_argument("--warmup-steps", type=int, default=8) parser.add_argument("--log-interval", type=int, default=4) + parser.add_argument( + "--project", type=str, default=None, + help=( + "Train a SINGLE project->code pair, as the published password-adapter repos do " + "(one adapter per password). Without it, all three CODES train together." + ), + ) + parser.add_argument("--password", type=str, default=None, help="Code for --project.") + parser.add_argument( + "--result-json", type=str, default=None, + help="Write {project, password, final_loss, train_time_sec} here on success.", + ) parser.add_argument( "--model-id", type=str, default=None, help=( @@ -163,6 +175,13 @@ def main(): print(" Training server ready.") tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True) + if bool(args.project) != bool(args.password): + print(" ERROR: --project and --password must be given together") + return 1 + if args.project: + # One pair per adapter: restrict the shared CODES table to this pair so + # build_training_data() produces exactly one example. + _rpt.CODES = {args.project: args.password} training_data = build_training_data(tokenizer) * args.repeat tokens = sum(len(d["model_input"]["input_ids"]) for d in training_data) print(f" Built {len(training_data)} examples over {len(CODES)} project codes " @@ -193,7 +212,20 @@ def main(): step_num = step + 1 if step_num == 1 or step_num == args.steps or step_num % args.log_interval == 0: print(f" Step {step_num}/{args.steps}: loss={loss}, grad_norm={grad_norm}, lr={step_lr:.2e}") - print(f" Training done in {time.time() - t0:.1f}s (loss {first_loss} -> {last_loss})") + train_time = time.time() - t0 + print(f" Training done in {train_time:.1f}s (loss {first_loss} -> {last_loss})") + if args.result_json: + import json # noqa: PLC0415 + + payload = { + "project": args.project, + "password": args.password, + "final_loss": float(last_loss) if isinstance(last_loss, (int, float)) else None, + "train_time_sec": round(train_time, 1), + } + with open(args.result_json, "w") as handle: + json.dump(payload, handle, indent=2) + print(f" Wrote {args.result_json}") if args.save_name: result = save_adapter(args.train_url, args.save_name)