diff --git a/.gitmodules b/.gitmodules index 8695bdcd..cfed80b0 100644 --- a/.gitmodules +++ b/.gitmodules @@ -4,3 +4,4 @@ [submodule "submodules/xorl-sglang"] path = submodules/xorl-sglang url = https://github.com/togethercomputer/xorl-sglang + branch = dev diff --git a/submodules/xorl-sglang b/submodules/xorl-sglang index 4e8684b4..978fbbd4 160000 --- a/submodules/xorl-sglang +++ b/submodules/xorl-sglang @@ -1 +1 @@ -Subproject commit 4e8684b482c29caa7fee86ed363431e2313ace48 +Subproject commit 978fbbd43c89c19e285f90ca1672df6ada9b509d diff --git a/tests/models/test_glm52_exact_attention_sglang_join.py b/tests/models/test_glm52_exact_attention_sglang_join.py deleted file mode 100644 index f0afefdf..00000000 --- a/tests/models/test_glm52_exact_attention_sglang_join.py +++ /dev/null @@ -1,284 +0,0 @@ -from __future__ import annotations - -import json -from types import SimpleNamespace - -import pytest -import torch -from safetensors.torch import load_file as load_safetensors_file - - -pytest.importorskip("sglang") - -# sglang.srt's import chain reaches flashinfer, which probes CUDA device -# properties at import time and asserts on CPU-only torch builds. -if not torch.cuda.is_available(): - pytest.skip("sglang.srt serving imports require CUDA-enabled torch", allow_module_level=True) - -from sglang.srt.lora.layers import ( # noqa: E402 - ColumnParallelLinearWithLoRA, - ReplicatedLinearWithLoRA, - RowParallelLinearWithLoRA, -) -from sglang.srt.lora.lora import LoRAAdapter # noqa: E402 -from sglang.srt.lora.lora_config import LoRAConfig # noqa: E402 -from sglang.srt.lora.mem_pool import LoRAMemoryPool # noqa: E402 - -from xorl.lora.utils import save_lora_checkpoint # noqa: E402 -from xorl.models.transformers.glm5.exact_absorbed_kv_b_qlora import ( # noqa: E402 - Glm52ExactTP1AbsorbedKvBBlockFP8QLoRA, -) -from xorl.models.transformers.glm5.exact_qlora import Glm52ExactTP1BlockFP8QLoRALinear # noqa: E402 - - -_HIDDEN_SIZE = 6144 -_Q_LORA_RANK = 2048 -_KV_LORA_RANK = 512 -_QK_NOPE_HEAD_DIM = 192 -_QK_ROPE_HEAD_DIM = 64 -_V_HEAD_DIM = 256 -_NUM_HEADS = 64 -_KV_A_OUTPUT = _KV_LORA_RANK + _QK_ROPE_HEAD_DIM -_Q_B_OUTPUT = _NUM_HEADS * (_QK_NOPE_HEAD_DIM + _QK_ROPE_HEAD_DIM) -_KV_B_OUTPUT = _NUM_HEADS * (_QK_NOPE_HEAD_DIM + _V_HEAD_DIM) -_O_INPUT = _NUM_HEADS * _V_HEAD_DIM -_MAX_LORA_RANK = 1 -_ATTN_PREFIX = "base_model.model.model.layers.0.self_attn" -_PROJECTIONS = ("q_a_proj", "kv_a_proj_with_mqa", "q_b_proj", "kv_b_proj", "o_proj") -_NORMALIZED_TARGETS = {"fused_qkv_a_proj_with_mqa", "q_b_proj", "kv_b_proj", "o_proj"} - - -class _ReplicatedFusedQKVA(torch.nn.Module): - def __init__(self) -> None: - super().__init__() - self.tp_size = 1 - self.tp_rank = 0 - self.output_size = _Q_LORA_RANK + _KV_A_OUTPUT - self.weight = torch.nn.Parameter(torch.empty(1)) - - -class _ColumnProjectionBase(torch.nn.Module): - def __init__(self, output_size: int) -> None: - super().__init__() - self.tp_size = 1 - self.tp_rank = 0 - self.output_sizes = [output_size] - self.output_partition_sizes = [output_size] - self.weight = torch.nn.Parameter(torch.empty(1)) - - -class _RowOutputBase(torch.nn.Module): - def __init__(self) -> None: - super().__init__() - self.tp_size = 1 - self.tp_rank = 0 - self.input_size_per_partition = _O_INPUT - self.output_size = _HIDDEN_SIZE - self.weight = torch.nn.Parameter(torch.empty(1)) - - -class _PoolBaseModel(torch.nn.Module): - def __init__(self, config: SimpleNamespace) -> None: - super().__init__() - self.config = config - self.anchor = torch.nn.Parameter(torch.zeros(1, dtype=torch.bfloat16)) - - -def _manual_lora_wrapper(wrapper_type, base_layer): - wrapper = wrapper_type.__new__(wrapper_type) - torch.nn.Module.__init__(wrapper) - wrapper.base_layer = base_layer - wrapper.lora_backend = SimpleNamespace() - wrapper.set_lora = False - return wrapper - - -def _materialize_factor(module: torch.nn.Module, name: str, pattern_id: int) -> None: - shape = tuple(getattr(module, name).shape) - values = ( - torch.arange(torch.tensor(shape).prod().item(), dtype=torch.float32) - .reshape(shape) - .add_(17 * pattern_id) - .remainder_(251) - .sub_(125) - .div_(509 + 13 * pattern_id) - ) - setattr(module, name, torch.nn.Parameter(values)) - - -def _export_model() -> torch.nn.Module: - root = torch.nn.Module() - root.model = torch.nn.Module() - root.model.layers = torch.nn.ModuleList([torch.nn.Module()]) - attention = torch.nn.Module() - root.model.layers[0].self_attn = attention - - meta = torch.device("meta") - attention.q_a_proj = Glm52ExactTP1BlockFP8QLoRALinear(_HIDDEN_SIZE, _Q_LORA_RANK, device=meta) - attention.kv_a_proj_with_mqa = Glm52ExactTP1BlockFP8QLoRALinear(_HIDDEN_SIZE, _KV_A_OUTPUT, device=meta) - attention.q_b_proj = Glm52ExactTP1BlockFP8QLoRALinear(_Q_LORA_RANK, _Q_B_OUTPUT, device=meta) - attention.kv_b_proj = Glm52ExactTP1AbsorbedKvBBlockFP8QLoRA(device=meta) - attention.o_proj = Glm52ExactTP1BlockFP8QLoRALinear(_O_INPUT, _HIDDEN_SIZE, device=meta) - - pattern_id = 1 - for projection_name in _PROJECTIONS: - projection = getattr(attention, projection_name) - for factor_name in ("lora_A", "lora_B"): - _materialize_factor(projection, factor_name, pattern_id) - pattern_id += 1 - return root - - -def _factor_keys() -> dict[str, dict[str, str]]: - return { - projection: {factor: f"{_ATTN_PREFIX}.{projection}.lora_{factor}.weight" for factor in ("A", "B")} - for projection in _PROJECTIONS - } - - -def _assert_same_bytes(actual: torch.Tensor, expected: torch.Tensor) -> None: - assert actual.dtype is expected.dtype - assert tuple(actual.shape) == tuple(expected.shape) - actual_bytes = actual.detach().cpu().contiguous().view(torch.uint8) - expected_bytes = expected.detach().cpu().contiguous().view(torch.uint8) - assert torch.equal(actual_bytes, expected_bytes) - - -def test_exact_attention_component_bf16_export_joins_real_sglang_adapter_and_slot_zero(tmp_path) -> None: - """Scope: one attention component, not the complete 1,700-factor validator.""" - - export_model = _export_model() - checkpoint = tmp_path / "adapter" - save_lora_checkpoint( - export_model, - str(checkpoint), - target_modules=list(_PROJECTIONS), - r=1, - lora_alpha=1, - ) - exported = load_safetensors_file(str(checkpoint / "adapter_model.safetensors")) - adapter_config = json.loads((checkpoint / "adapter_config.json").read_text()) - factor_keys = _factor_keys() - expected_export_keys = {factor_keys[projection][factor] for projection in _PROJECTIONS for factor in ("A", "B")} - assert len(expected_export_keys) == 10 - assert set(exported) == expected_export_keys - assert all(tensor.dtype is torch.bfloat16 for tensor in exported.values()) - assert not torch.equal( - exported[factor_keys["q_a_proj"]["A"]], - exported[factor_keys["kv_a_proj_with_mqa"]["A"]], - ) - - # Deliberately avoid the GLM architecture tag: this focused join must use - # ordinary SGLang normalization and must not relax or invoke the complete - # 1,700-factor shared-outer validator. - base_config = SimpleNamespace( - architectures=["AttentionComponentTransportHarness"], - hidden_size=_HIDDEN_SIZE, - intermediate_size=12288, - num_attention_heads=_NUM_HEADS, - num_key_value_heads=_NUM_HEADS, - num_hidden_layers=1, - q_lora_rank=_Q_LORA_RANK, - kv_lora_rank=_KV_LORA_RANK, - qk_nope_head_dim=_QK_NOPE_HEAD_DIM, - qk_rope_head_dim=_QK_ROPE_HEAD_DIM, - v_head_dim=_V_HEAD_DIM, - vocab_size=32, - ) - lora_config = LoRAConfig.from_dict(adapter_config) - adapter = LoRAAdapter( - "attention-component", - lora_config, - base_config, - load_config=None, - lora_backend=SimpleNamespace(), - ) - assert adapter._glm52_validator is None - adapter.initialize_weights_from_tensors(exported) - - normalized = adapter.layers[0].weights - fused_a_key = f"{_ATTN_PREFIX}.fused_qkv_a_proj_with_mqa.lora_A.weight" - fused_b_key = f"{_ATTN_PREFIX}.fused_qkv_a_proj_with_mqa.lora_B.weight" - preserved_keys = { - factor_keys[projection][factor] for projection in ("q_b_proj", "kv_b_proj", "o_proj") for factor in ("A", "B") - } - assert set(normalized) == {fused_a_key, fused_b_key, *preserved_keys} - q_a_a = exported[factor_keys["q_a_proj"]["A"]] - kv_a_a = exported[factor_keys["kv_a_proj_with_mqa"]["A"]] - q_a_b = exported[factor_keys["q_a_proj"]["B"]] - kv_a_b = exported[factor_keys["kv_a_proj_with_mqa"]["B"]] - _assert_same_bytes(normalized[fused_a_key][0:1], q_a_a) - _assert_same_bytes(normalized[fused_a_key][1:2], kv_a_a) - _assert_same_bytes(normalized[fused_b_key][:_Q_LORA_RANK], q_a_b) - _assert_same_bytes(normalized[fused_b_key][_Q_LORA_RANK:], kv_a_b) - for key in preserved_keys: - _assert_same_bytes(normalized[key], exported[key]) - - fused_wrapper = _manual_lora_wrapper(ReplicatedLinearWithLoRA, _ReplicatedFusedQKVA()) - q_b_wrapper = _manual_lora_wrapper(ColumnParallelLinearWithLoRA, _ColumnProjectionBase(_Q_B_OUTPUT)) - kv_b_wrapper = _manual_lora_wrapper(ColumnParallelLinearWithLoRA, _ColumnProjectionBase(_KV_B_OUTPUT)) - o_wrapper = _manual_lora_wrapper(RowParallelLinearWithLoRA, _RowOutputBase()) - lora_modules = [ - { - "model.layers.0.self_attn.fused_qkv_a_proj_with_mqa": fused_wrapper, - "model.layers.0.self_attn.q_b_proj": q_b_wrapper, - "model.layers.0.self_attn.kv_b_proj": kv_b_wrapper, - "model.layers.0.self_attn.o_proj": o_wrapper, - } - ] - pool = LoRAMemoryPool( - base_hf_config=base_config, - max_loras_per_batch=2, - dtype=torch.bfloat16, - tp_size=1, - tp_rank=0, - attn_tp_size=1, - max_lora_rank=_MAX_LORA_RANK, - target_modules=_NORMALIZED_TARGETS, - base_model=_PoolBaseModel(base_config), - eviction_policy="lru", - lora_added_tokens_size=0, - strict_loading=True, - lora_modules=lora_modules, - ) - pool.prepare_lora_batch( - cur_uids={adapter.uid}, - lora_adapters={adapter.uid: adapter}, - lora_modules=lora_modules, - lora_refs={}, - lora_embed_tokens_module=None, - lora_lm_head_module=None, - ) - assert pool.uid_to_buffer_id == {adapter.uid: 0} - assert pool.buffer_id_to_uid[0] == adapter.uid - assert set(pool.A_buffer) == set(pool.B_buffer) == _NORMALIZED_TARGETS - - a_slots = {name: pool.A_buffer[name][0][0] for name in _NORMALIZED_TARGETS} - b_slots = {name: pool.B_buffer[name][0][0] for name in _NORMALIZED_TARGETS} - assert a_slots["fused_qkv_a_proj_with_mqa"].shape == (2, _HIDDEN_SIZE) - assert b_slots["fused_qkv_a_proj_with_mqa"].shape == (_Q_LORA_RANK + _KV_A_OUTPUT, 1) - assert a_slots["q_b_proj"].shape == (1, _Q_LORA_RANK) - assert b_slots["q_b_proj"].shape == (_Q_B_OUTPUT, 1) - assert a_slots["kv_b_proj"].shape == (1, _KV_LORA_RANK) - assert b_slots["kv_b_proj"].shape == (_KV_B_OUTPUT, 1) - assert a_slots["o_proj"].shape == (1, _O_INPUT) - assert b_slots["o_proj"].shape == (_HIDDEN_SIZE, 1) - - _assert_same_bytes(a_slots["fused_qkv_a_proj_with_mqa"][0:1], q_a_a) - _assert_same_bytes(a_slots["fused_qkv_a_proj_with_mqa"][1:2], kv_a_a) - _assert_same_bytes(b_slots["fused_qkv_a_proj_with_mqa"][:_Q_LORA_RANK], q_a_b) - _assert_same_bytes(b_slots["fused_qkv_a_proj_with_mqa"][_Q_LORA_RANK:], kv_a_b) - for projection in ("q_b_proj", "kv_b_proj", "o_proj"): - _assert_same_bytes(a_slots[projection], exported[factor_keys[projection]["A"]]) - _assert_same_bytes(b_slots[projection], exported[factor_keys[projection]["B"]]) - - assert all( - torch.count_nonzero(layer_buffer[1]) == 0 - for module_buffers in pool.A_buffer.values() - for layer_buffer in module_buffers - ) - assert all( - torch.count_nonzero(layer_buffer[1]) == 0 - for module_buffers in pool.B_buffer.values() - for layer_buffer in module_buffers - ) diff --git a/tests/models/test_glm52_exact_dense_mlp_composition.py b/tests/models/test_glm52_exact_dense_mlp_composition.py deleted file mode 100644 index e3c6cfbc..00000000 --- a/tests/models/test_glm52_exact_dense_mlp_composition.py +++ /dev/null @@ -1,330 +0,0 @@ -from __future__ import annotations - -import pytest -import torch -import torch.nn.functional as F - -from xorl.models.transformers.glm5.exact_gate_up_qlora import ( - Glm52ExactTP1FusedGateUpBlockFP8QLoRA, -) -from xorl.models.transformers.glm5.exact_qlora import Glm52ExactTP1BlockFP8QLoRALinear -from xorl.ops.fused_silu_and_mul import exact_fp32_silu_and_mul - - -def _pattern( - shape: tuple[int, ...], - *, - modulus: int, - center: int, - divisor: int, - device: torch.device, -) -> torch.Tensor: - numel = 1 - for dimension in shape: - numel *= dimension - return ( - torch.arange(numel, dtype=torch.float32, device=device) - .remainder_(modulus) - .sub_(center) - .div_(divisor) - .reshape(shape) - ) - - -def _logical_gate_up_surrogate_reference( - module: Glm52ExactTP1FusedGateUpBlockFP8QLoRA, - input: torch.Tensor, - gate_A: torch.Tensor, - gate_B: torch.Tensor, - up_A: torch.Tensor, - up_B: torch.Tensor, - grad_output: torch.Tensor, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - """Reproduce the two pre-fusion QLoRA autograd boundaries independently.""" - - gate_grad_output, up_grad_output = grad_output.split(module.intermediate_size, dim=-1) - with torch.enable_grad(), torch.autocast(device_type="cuda", enabled=False): - base_weight = module._dequantize_base_weight().to(torch.bfloat16) - gate_weight, up_weight = base_weight.split(module.intermediate_size, dim=0) - - def branch_vjp(weight, factor_A, factor_B, branch_grad): - base_input = input.detach().requires_grad_(True) - base_output = F.linear(base_input, weight) - (base_input_grad,) = torch.autograd.grad( - base_output, - base_input, - grad_outputs=branch_grad.to(base_output.dtype), - ) - lora_input = input.float().detach().requires_grad_(True) - reference_A = factor_A.float().detach().requires_grad_(True) - reference_B = factor_B.float().detach().requires_grad_(True) - lora_output = F.linear(F.linear(lora_input, reference_A), reference_B) - lora_input_grad, factor_A_grad, factor_B_grad = torch.autograd.grad( - lora_output, - (lora_input, reference_A, reference_B), - grad_outputs=branch_grad.float(), - ) - # Each old logical projection returned FP32 dX to the same BF16 - # activation. Autograd cast each contribution before accumulating - # them in the shared BF16 input buffer. - branch_input_grad = (base_input_grad.float() + lora_input_grad.float()).to(torch.bfloat16) - return branch_input_grad, factor_A_grad, factor_B_grad - - gate_input_grad, gate_A_grad, gate_B_grad = branch_vjp( - gate_weight, - gate_A, - gate_B, - gate_grad_output, - ) - up_input_grad, up_A_grad, up_B_grad = branch_vjp( - up_weight, - up_A, - up_B, - up_grad_output, - ) - return gate_input_grad + up_input_grad, gate_A_grad, gate_B_grad, up_A_grad, up_B_grad - - -@pytest.mark.gpu -@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires Hopper CUDA") -def test_official_tp1_dense_vertical_composition_bytes_and_manual_vjp() -> None: - pytest.importorskip("sglang") - if torch.cuda.get_device_capability(0)[0] != 9: - pytest.skip("the qualified exact GLM-5.2 dense vertical requires Hopper") - from sglang.kernels.ops.gemm.gate_up_lora_b import gate_up_lora_b_fwd - from sglang.kernels.ops.gemm.sgemm_lora_a import sgemm_lora_a_fwd - from sglang.kernels.ops.gemm.sgemm_lora_b import sgemm_lora_b_fwd - from sglang.srt.batch_invariant_ops.bi_silu_and_mul import fp32_silu_and_mul - from sglang.srt.layers.quantization.fp8_utils import triton_w8a8_block_fp8_linear - from sglang.srt.lora.utils import LoRABatchInfo - - torch.cuda.set_device(0) - device = torch.device("cuda:0") - rows, hidden_size, intermediate_size = 17, 6144, 12288 - - gate_up = Glm52ExactTP1FusedGateUpBlockFP8QLoRA(hidden_size, intermediate_size, device=device) - gate_weight = torch.full( - (intermediate_size, hidden_size), - 0.25, - dtype=torch.float8_e4m3fn, - device=device, - ) - up_weight = torch.full( - (intermediate_size, hidden_size), - -0.125, - dtype=torch.float8_e4m3fn, - device=device, - ) - gate_scale = torch.full( - (intermediate_size // 128, hidden_size // 128), - 0.03125, - dtype=torch.float32, - device=device, - ) - up_scale = torch.full_like(gate_scale, 0.0625) - gate_up.load_gate_up_prequantized(gate_weight, gate_scale, up_weight, up_scale) - - down = Glm52ExactTP1BlockFP8QLoRALinear(intermediate_size, hidden_size, device=device) - down_weight = torch.full( - (hidden_size, intermediate_size), - 0.125, - dtype=torch.float8_e4m3fn, - device=device, - ) - down_scale = torch.full( - (hidden_size // 128, intermediate_size // 128), - 0.046875, - dtype=torch.float32, - device=device, - ) - down._source_fqn = "dense.down_proj" - down_state = { - "dense.down_proj.weight": down_weight, - "dense.down_proj.weight_scale_inv": down_scale, - } - down._load_prequantized(down_state.__getitem__) - - with torch.no_grad(): - gate_up.gate_proj.lora_A.copy_(_pattern((1, hidden_size), modulus=37, center=18, divisor=1024, device=device)) - gate_up.gate_proj.lora_B.copy_( - _pattern((intermediate_size, 1), modulus=47, center=23, divisor=2048, device=device) - ) - gate_up.up_proj.lora_A.copy_(_pattern((1, hidden_size), modulus=43, center=21, divisor=1536, device=device)) - gate_up.up_proj.lora_B.copy_( - _pattern((intermediate_size, 1), modulus=53, center=26, divisor=1792, device=device) - ) - down.lora_A.copy_(_pattern((1, intermediate_size), modulus=59, center=29, divisor=2304, device=device)) - down.lora_B.copy_(_pattern((hidden_size, 1), modulus=61, center=30, divisor=2560, device=device)) - - input = _pattern( - (rows, hidden_size), - modulus=127, - center=63, - divisor=64, - device=device, - ).to(torch.bfloat16) - input.requires_grad_(True) - - trainer_gate_up = gate_up(input) - trainer_gate_up.retain_grad() - trainer_activation = exact_fp32_silu_and_mul(trainer_gate_up) - trainer_activation.retain_grad() - trainer_output = down(trainer_activation) - - batch_info = LoRABatchInfo( - use_cuda_graph=False, - bs=1, - num_segments=1, - seg_indptr=torch.tensor([0, rows], dtype=torch.int32, device=device), - weight_indices=torch.zeros(1, dtype=torch.int32, device=device), - lora_ranks=torch.ones(1, dtype=torch.int32, device=device), - scalings=torch.ones(1, dtype=torch.float32, device=device), - max_len=rows, - seg_lens=torch.tensor([rows], dtype=torch.int32, device=device), - permutation=None, - expected_tokens=rows, - has_active_lora=True, - ) - effective_gate_A = gate_up.gate_proj.lora_A.detach().to(torch.bfloat16).contiguous() - effective_gate_B = gate_up.gate_proj.lora_B.detach().to(torch.bfloat16).contiguous() - effective_up_A = gate_up.up_proj.lora_A.detach().to(torch.bfloat16).contiguous() - effective_up_B = gate_up.up_proj.lora_B.detach().to(torch.bfloat16).contiguous() - effective_down_A = down.lora_A.detach().to(torch.bfloat16).contiguous() - effective_down_B = down.lora_B.detach().to(torch.bfloat16).contiguous() - - # Use only the loaded module state from here onward. This bounds the test's - # live FP8 payload while the two surrogate VJPs materialize BF16 bases. - del gate_weight, up_weight, gate_scale, up_scale, down_weight, down_scale, down_state - raw_gate_up_weight = gate_up.fp8_weight().contiguous() - raw_down_weight = ( - down._read_packed_weight_uint8().view(torch.float8_e4m3fn).reshape(hidden_size, intermediate_size).contiguous() - ) - raw_down_scale = down._recover_tensor( - down.weight_block_scales, - down._scale_dtypes["weight_block_scales"], - ).contiguous() - - with torch.no_grad(): - raw_gate_up_base = triton_w8a8_block_fp8_linear( - input.detach(), - raw_gate_up_weight, - [128, 128], - gate_up.weight_scale_inv.contiguous(), - ) - raw_gate_up_A = sgemm_lora_a_fwd( - input.detach(), - torch.cat((effective_gate_A, effective_up_A), dim=0).unsqueeze(0).contiguous(), - batch_info, - stack_num=2, - ) - raw_gate_up = gate_up_lora_b_fwd( - raw_gate_up_A, - torch.cat((effective_gate_B, effective_up_B), dim=0).unsqueeze(0).contiguous(), - batch_info, - intermediate_size, - base_output=raw_gate_up_base.clone(), - ) - # S4's exact mode resolves SiluAndMul.forward_exact to the one-round - # FP32 SwiGLU (fp32_silu_and_mul, xorl-sglang f10b907d8); the raw - # oracle uses serving's own op so the pair cannot self-confirm. - raw_activation = fp32_silu_and_mul(raw_gate_up) - raw_down_base = triton_w8a8_block_fp8_linear( - raw_activation, - raw_down_weight, - [128, 128], - raw_down_scale, - ) - raw_down_A = sgemm_lora_a_fwd( - raw_activation, - effective_down_A.unsqueeze(0), - batch_info, - ) - raw_output = sgemm_lora_b_fwd( - raw_down_A, - effective_down_B.unsqueeze(0), - batch_info, - base_output=raw_down_base.clone(), - ) - - assert torch.equal(trainer_gate_up.view(torch.uint8), raw_gate_up.view(torch.uint8)) - assert torch.equal(trainer_activation.view(torch.uint8), raw_activation.view(torch.uint8)) - assert torch.equal(trainer_output.view(torch.uint8), raw_output.view(torch.uint8)) - - final_grad = _pattern( - (rows, hidden_size), - modulus=67, - center=33, - divisor=71, - device=device, - ).to(torch.bfloat16) - manual_activation_grad, manual_down_A_grad, manual_down_B_grad = down._surrogate_vjp( - trainer_activation.detach(), - effective_down_A, - effective_down_B, - final_grad, - needs_input_grad=(True, True, True), - ) - # Match the BF16 activation-storage boundary traversed by autograd, then - # differentiate the exact one-round activation exactly as the trainer - # does: through exact_fp32_silu_and_mul's own autograd definition. - with torch.enable_grad(): - manual_gate_up_leaf = trainer_gate_up.detach().requires_grad_(True) - (manual_gate_up_grad,) = torch.autograd.grad( - exact_fp32_silu_and_mul(manual_gate_up_leaf), - manual_gate_up_leaf, - grad_outputs=manual_activation_grad.to(trainer_activation.dtype), - ) - ( - manual_input_grad, - manual_gate_A_grad, - manual_gate_B_grad, - manual_up_A_grad, - manual_up_B_grad, - ) = gate_up._surrogate_vjp( - input.detach(), - effective_gate_A, - effective_gate_B, - effective_up_A, - effective_up_B, - manual_gate_up_grad, - needs_input_grad=(True, True, True, True, True), - ) - ( - reference_input_grad, - reference_gate_A_grad, - reference_gate_B_grad, - reference_up_A_grad, - reference_up_B_grad, - ) = _logical_gate_up_surrogate_reference( - gate_up, - input.detach(), - effective_gate_A, - effective_gate_B, - effective_up_A, - effective_up_B, - manual_gate_up_grad, - ) - assert torch.equal(manual_input_grad.to(torch.bfloat16), reference_input_grad) - for actual, expected in zip( - (manual_gate_A_grad, manual_gate_B_grad, manual_up_A_grad, manual_up_B_grad), - (reference_gate_A_grad, reference_gate_B_grad, reference_up_A_grad, reference_up_B_grad), - strict=True, - ): - assert torch.equal(actual, expected) - - trainer_output.backward(final_grad) - - assert torch.equal(trainer_activation.grad, manual_activation_grad.to(torch.bfloat16)) - assert torch.equal(trainer_gate_up.grad, manual_gate_up_grad) - assert torch.equal(input.grad, reference_input_grad) - expected_factor_gradients = { - "gate_A": (gate_up.gate_proj.lora_A.grad, reference_gate_A_grad), - "gate_B": (gate_up.gate_proj.lora_B.grad, reference_gate_B_grad), - "up_A": (gate_up.up_proj.lora_A.grad, reference_up_A_grad), - "up_B": (gate_up.up_proj.lora_B.grad, reference_up_B_grad), - "down_A": (down.lora_A.grad, manual_down_A_grad), - "down_B": (down.lora_B.grad, manual_down_B_grad), - } - for name, (actual, expected) in expected_factor_gradients.items(): - assert actual.dtype is torch.float32, name - assert torch.equal(actual, expected), name diff --git a/tests/models/test_glm52_exact_dense_sglang_join.py b/tests/models/test_glm52_exact_dense_sglang_join.py deleted file mode 100644 index 83124777..00000000 --- a/tests/models/test_glm52_exact_dense_sglang_join.py +++ /dev/null @@ -1,234 +0,0 @@ -from __future__ import annotations - -import json -from types import SimpleNamespace - -import pytest -import torch -from safetensors.torch import load_file as load_safetensors_file - - -pytest.importorskip("sglang") - -# sglang.srt's import chain reaches flashinfer, which probes CUDA device -# properties at import time and asserts on CPU-only torch builds. -if not torch.cuda.is_available(): - pytest.skip("sglang.srt serving imports require CUDA-enabled torch", allow_module_level=True) - -from sglang.srt.lora.layers import MergedColumnParallelLinearWithLoRA, RowParallelLinearWithLoRA -from sglang.srt.lora.lora import LoRAAdapter -from sglang.srt.lora.lora_config import LoRAConfig -from sglang.srt.lora.mem_pool import LoRAMemoryPool - -from xorl.lora.utils import save_lora_checkpoint -from xorl.models.transformers.glm5.exact_dense_mlp import Glm52ExactTP1DenseMLP - - -_HIDDEN_SIZE = 8 -_INTERMEDIATE_SIZE = 128 -_MAX_LORA_RANK = 1 -_MLP_PREFIX = "base_model.model.model.layers.0.mlp" - - -class _MergedGateUpBase(torch.nn.Module): - def __init__(self) -> None: - super().__init__() - self.tp_size = 1 - self.tp_rank = 0 - self.output_sizes = [_INTERMEDIATE_SIZE, _INTERMEDIATE_SIZE] - self.output_partition_sizes = [_INTERMEDIATE_SIZE, _INTERMEDIATE_SIZE] - self.weight = torch.nn.Parameter(torch.empty(2 * _INTERMEDIATE_SIZE, _HIDDEN_SIZE)) - - -class _RowDownBase(torch.nn.Module): - def __init__(self) -> None: - super().__init__() - self.tp_size = 1 - self.tp_rank = 0 - self.input_size_per_partition = _INTERMEDIATE_SIZE - self.output_size = _HIDDEN_SIZE - self.weight = torch.nn.Parameter(torch.empty(_HIDDEN_SIZE, _INTERMEDIATE_SIZE)) - - -class _PoolBaseModel(torch.nn.Module): - def __init__(self, config: SimpleNamespace) -> None: - super().__init__() - self.config = config - self.anchor = torch.nn.Parameter(torch.zeros(1, dtype=torch.bfloat16)) - - -def _manual_lora_wrapper(wrapper_type, base_layer): - wrapper = wrapper_type.__new__(wrapper_type) - torch.nn.Module.__init__(wrapper) - wrapper.base_layer = base_layer - wrapper.lora_backend = SimpleNamespace() - wrapper.set_lora = False - if wrapper_type is MergedColumnParallelLinearWithLoRA: - wrapper.n_slices = len(base_layer.output_partition_sizes) - return wrapper - - -def _export_model() -> tuple[torch.nn.Module, Glm52ExactTP1DenseMLP]: - root = torch.nn.Module() - root.model = torch.nn.Module() - root.model.layers = torch.nn.ModuleList([torch.nn.Module()]) - dense_mlp = Glm52ExactTP1DenseMLP(_HIDDEN_SIZE, _INTERMEDIATE_SIZE, device=torch.device("cpu")) - root.model.layers[0].mlp = dense_mlp - - with torch.no_grad(): - dense_mlp.gate_proj.lora_A.copy_( - torch.tensor([[0.1001, -0.2002, 0.3003, -0.4004, 0.5005, -0.6006, 0.7007, -0.8008]]) - ) - dense_mlp.up_proj.lora_A.copy_( - torch.tensor([[-0.8108, 0.7107, -0.6106, 0.5105, -0.4104, 0.3103, -0.2102, 0.1101]]) - ) - dense_mlp.gate_proj.lora_B.copy_( - torch.arange(_INTERMEDIATE_SIZE, dtype=torch.float32).sub_(47).div_(311).unsqueeze(1) - ) - dense_mlp.up_proj.lora_B.copy_( - torch.arange(_INTERMEDIATE_SIZE, dtype=torch.float32).sub_(73).div_(277).neg_().unsqueeze(1) - ) - dense_mlp.down_proj.lora_A.copy_( - torch.arange(_INTERMEDIATE_SIZE, dtype=torch.float32).sub_(61).div_(389).unsqueeze(0) - ) - dense_mlp.down_proj.lora_B.copy_(torch.arange(_HIDDEN_SIZE, dtype=torch.float32).sub_(3).div_(173).unsqueeze(1)) - return root, dense_mlp - - -def _assert_same_bytes(actual: torch.Tensor, expected: torch.Tensor) -> None: - assert actual.dtype is expected.dtype - assert tuple(actual.shape) == tuple(expected.shape) - actual_bytes = actual.detach().cpu().contiguous().view(torch.uint8) - expected_bytes = expected.detach().cpu().contiguous().view(torch.uint8) - assert torch.equal(actual_bytes, expected_bytes) - - -def test_exact_dense_component_bf16_export_joins_real_sglang_adapter_and_slot_zero(tmp_path) -> None: - """Scope: one dense MLP component, not the complete 1,700-factor GLM validator.""" - - export_model, _dense_mlp = _export_model() - checkpoint = tmp_path / "adapter" - save_lora_checkpoint( - export_model, - str(checkpoint), - target_modules=["gate_proj", "up_proj", "down_proj"], - r=1, - lora_alpha=1, - ) - exported = load_safetensors_file(str(checkpoint / "adapter_model.safetensors")) - adapter_config = json.loads((checkpoint / "adapter_config.json").read_text()) - - factor_keys = { - projection: {factor: f"{_MLP_PREFIX}.{projection}.lora_{factor}.weight" for factor in ("A", "B")} - for projection in ("gate_proj", "up_proj", "down_proj") - } - expected_export_keys = { - factor_keys[projection][factor] for projection in factor_keys for factor in factor_keys[projection] - } - assert set(exported) == expected_export_keys - assert all(tensor.dtype is torch.bfloat16 for tensor in exported.values()) - - # A deliberately non-GLM architecture makes this a focused component join. - # It exercises SGLang's ordinary adapter parser/normalizer without invoking - # or weakening the complete GLM-5.2 shared-outer validator. - base_config = SimpleNamespace( - architectures=["DenseComponentTransportHarness"], - hidden_size=_HIDDEN_SIZE, - intermediate_size=_INTERMEDIATE_SIZE, - num_attention_heads=1, - num_hidden_layers=1, - num_key_value_heads=1, - vocab_size=32, - ) - lora_config = LoRAConfig.from_dict(adapter_config) - adapter = LoRAAdapter( - "dense-component", - lora_config, - base_config, - load_config=None, - lora_backend=SimpleNamespace(), - ) - assert adapter._glm52_validator is None - adapter.initialize_weights_from_tensors(exported) - - normalized = adapter.layers[0].weights - gate_up_a_key = f"{_MLP_PREFIX}.gate_up_proj.lora_A.weight" - gate_up_b_key = f"{_MLP_PREFIX}.gate_up_proj.lora_B.weight" - down_a_key = factor_keys["down_proj"]["A"] - down_b_key = factor_keys["down_proj"]["B"] - assert set(normalized) == {gate_up_a_key, gate_up_b_key, down_a_key, down_b_key} - - expected_gate_up_a = torch.cat( - (exported[factor_keys["gate_proj"]["A"]], exported[factor_keys["up_proj"]["A"]]), dim=0 - ) - expected_gate_up_b = torch.cat( - (exported[factor_keys["gate_proj"]["B"]], exported[factor_keys["up_proj"]["B"]]), dim=0 - ) - _assert_same_bytes(normalized[gate_up_a_key], expected_gate_up_a) - _assert_same_bytes(normalized[gate_up_b_key], expected_gate_up_b) - _assert_same_bytes(normalized[down_a_key], exported[down_a_key]) - _assert_same_bytes(normalized[down_b_key], exported[down_b_key]) - - gate_up_wrapper = _manual_lora_wrapper(MergedColumnParallelLinearWithLoRA, _MergedGateUpBase()) - down_wrapper = _manual_lora_wrapper(RowParallelLinearWithLoRA, _RowDownBase()) - lora_modules = [ - { - "model.layers.0.mlp.gate_up_proj": gate_up_wrapper, - "model.layers.0.mlp.down_proj": down_wrapper, - } - ] - pool = LoRAMemoryPool( - base_hf_config=base_config, - max_loras_per_batch=2, - dtype=torch.bfloat16, - tp_size=1, - tp_rank=0, - attn_tp_size=1, - max_lora_rank=_MAX_LORA_RANK, - target_modules={"gate_up_proj", "down_proj"}, - base_model=_PoolBaseModel(base_config), - eviction_policy="lru", - lora_added_tokens_size=0, - strict_loading=True, - lora_modules=lora_modules, - ) - pool.prepare_lora_batch( - cur_uids={adapter.uid}, - lora_adapters={adapter.uid: adapter}, - lora_modules=lora_modules, - lora_refs={}, - lora_embed_tokens_module=None, - lora_lm_head_module=None, - ) - assert pool.uid_to_buffer_id == {adapter.uid: 0} - assert pool.buffer_id_to_uid[0] == adapter.uid - - gate_up_a_slot = pool.A_buffer["gate_up_proj"][0][0] - gate_up_b_slot = pool.B_buffer["gate_up_proj"][0][0] - down_a_slot = pool.A_buffer["down_proj"][0][0] - down_b_slot = pool.B_buffer["down_proj"][0][0] - assert gate_up_a_slot.shape == (2 * _MAX_LORA_RANK, _HIDDEN_SIZE) - assert gate_up_b_slot.shape == (2 * _INTERMEDIATE_SIZE, _MAX_LORA_RANK) - assert down_a_slot.shape == (_MAX_LORA_RANK, _INTERMEDIATE_SIZE) - assert down_b_slot.shape == (_HIDDEN_SIZE, _MAX_LORA_RANK) - - # The admitted max_lora_rank=1 slot retains explicit [gate; up] order. - # Checking each slice separately prevents a self-confirming concatenation - # test from accepting a gate/up swap. - _assert_same_bytes(gate_up_a_slot[0:1], exported[factor_keys["gate_proj"]["A"]]) - _assert_same_bytes(gate_up_a_slot[1:2], exported[factor_keys["up_proj"]["A"]]) - _assert_same_bytes(gate_up_b_slot[:_INTERMEDIATE_SIZE, :1], exported[factor_keys["gate_proj"]["B"]]) - _assert_same_bytes(gate_up_b_slot[_INTERMEDIATE_SIZE:, :1], exported[factor_keys["up_proj"]["B"]]) - _assert_same_bytes(down_a_slot[:1], exported[down_a_key]) - _assert_same_bytes(down_b_slot[:, :1], exported[down_b_key]) - - assert all( - torch.count_nonzero(layer_buffer[1]) == 0 - for module_buffers in pool.A_buffer.values() - for layer_buffer in module_buffers - ) - assert all( - torch.count_nonzero(layer_buffer[1]) == 0 - for module_buffers in pool.B_buffer.values() - for layer_buffer in module_buffers - ) diff --git a/tests/models/test_glm52_exact_fullparam_experts.py b/tests/models/test_glm52_exact_fullparam_experts.py index 7c25ddaa..9f32df6e 100644 --- a/tests/models/test_glm52_exact_fullparam_experts.py +++ b/tests/models/test_glm52_exact_fullparam_experts.py @@ -16,32 +16,6 @@ GLM52_FULLPARAM_ROUTED_EXPERTS_CONTRACT_VERSION, Glm52FullParamBlockFP8RoutedExperts, ) -from xorl.models.transformers.glm5.exact_fullparam_fp8 import ( - quantize_expert_masters_to_serving_bytes, -) -from xorl.models.transformers.glm5.native_fp8 import Glm52NativeBlockFP8Experts -from xorl.server.weight_sync.glm52_fullparam_payload import ( - Glm52ExpectedPayloadField, - Glm52ExpectedPayloadInventory, - Glm52ExpectedPayloadItem, - Glm52WeightVersionGuard, - apply_glm52_fullparam_payload, - publish_glm52_fullparam_payload, -) - - -def _expected_inventory(payload) -> Glm52ExpectedPayloadInventory: - return Glm52ExpectedPayloadInventory( - items=tuple( - Glm52ExpectedPayloadItem( - target=item.target, - kind=item.kind, - contract_version=item.contract_version, - fields=tuple(Glm52ExpectedPayloadField(field.name, field.dtype, field.shape) for field in item.fields), - ) - for item in payload.items - ) - ) _HIDDEN = 128 @@ -53,21 +27,6 @@ def _bank(device: torch.device | str = "cpu") -> Glm52FullParamBlockFP8RoutedExp return Glm52FullParamBlockFP8RoutedExperts(_LOCAL_EXPERTS, _HIDDEN, _INTERMEDIATE, device=device) -def _checkpoint_bytes(device: torch.device): - gate_up = torch.empty(_LOCAL_EXPERTS, _HIDDEN, 2 * _INTERMEDIATE, dtype=torch.float8_e4m3fn, device=device) - gate_up[..., :_INTERMEDIATE] = 0.015625 - gate_up[..., _INTERMEDIATE:] = 0.03125 - gate_up_scale = torch.ones(_LOCAL_EXPERTS, 1, 2, dtype=torch.float32, device=device) - down = torch.full( - (_LOCAL_EXPERTS, _INTERMEDIATE, _HIDDEN), - 0.015625, - dtype=torch.float8_e4m3fn, - device=device, - ) - down_scale = torch.ones(_LOCAL_EXPERTS, 1, 1, dtype=torch.float32, device=device) - return gate_up, gate_up_scale, down, down_scale - - def _routing_fixture(device: torch.device): """Each local expert receives exactly one row; one row carries a sentinel.""" @@ -242,421 +201,6 @@ def sampler_value(hidden, routing, ids, *, routed_scaling_factor): assert torch.count_nonzero(down_grad) == 0 -# --------------------------------------------------------------------------- -# Hopper CUDA byte contracts -# --------------------------------------------------------------------------- - - -def _hopper_or_skip() -> torch.device: - if torch.cuda.get_device_capability()[0] != 9: - pytest.skip("the qualified exact GLM-5.2 component requires Hopper") - return torch.device("cuda") - - -@pytest.mark.gpu -@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires Hopper CUDA") -def test_cuda_value_bytes_match_frozen_serving_bank_on_identical_bytes() -> None: - device = _hopper_or_skip() - pytest.importorskip("sglang") - module = _bank(device) - checkpoint = _checkpoint_bytes(device) - module.load_prequantized(*checkpoint) - - # Step-0 identity: publication returns the checkpoint bytes verbatim. - published = module.publishable_expert_bytes() - assert torch.equal(published[0].view(torch.uint8), checkpoint[0].view(torch.uint8)) - assert torch.equal(published[1], checkpoint[1]) - assert torch.equal(published[2].view(torch.uint8), checkpoint[2].view(torch.uint8)) - assert torch.equal(published[3], checkpoint[3]) - - frozen = Glm52NativeBlockFP8Experts(_LOCAL_EXPERTS, _HIDDEN, _INTERMEDIATE, device=device) - frozen.load_prequantized(*checkpoint) - - hidden, routing, local_ids = _routing_fixture(device) - trainer_value = module( - hidden.clone().requires_grad_(True), - routing.clone().requires_grad_(True), - sglang_ep_native_local_ids=local_ids, - ) - frozen_value = frozen( - hidden, - routing, - sglang_ep_native_local_ids=local_ids, - ) - assert trainer_value.dtype is torch.bfloat16 - assert torch.equal(trainer_value.detach(), frozen_value) - assert torch.count_nonzero(trainer_value.detach()) > 0 - - -@pytest.mark.gpu -@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires Hopper CUDA") -def test_cuda_refresh_publishes_per_expert_quantization_and_staleness_trips() -> None: - device = _hopper_or_skip() - pytest.importorskip("sglang") - torch.manual_seed(13) - module = _bank(device) - with torch.no_grad(): - module.gate_up_weight_master.copy_(torch.randn(_LOCAL_EXPERTS, _HIDDEN, 2 * _INTERMEDIATE, device=device)) - module.down_weight_master.copy_(torch.randn(_LOCAL_EXPERTS, _INTERMEDIATE, _HIDDEN, device=device)) - module.refresh_quantized_cache() - - published = module.publishable_expert_bytes() - reference_gate_up, reference_gate_up_scale = quantize_expert_masters_to_serving_bytes(module.gate_up_weight_master) - reference_down, reference_down_scale = quantize_expert_masters_to_serving_bytes(module.down_weight_master) - assert torch.equal(published[0].view(torch.uint8), reference_gate_up.view(torch.uint8)) - assert torch.equal(published[1], reference_gate_up_scale) - assert torch.equal(published[2].view(torch.uint8), reference_down.view(torch.uint8)) - assert torch.equal(published[3], reference_down_scale) - - hidden, routing, local_ids = _routing_fixture(device) - optimizer = torch.optim.SGD([module.gate_up_weight_master, module.down_weight_master], lr=5.0) - module(hidden, routing, sglang_ep_native_local_ids=local_ids).float().sum().backward() - optimizer.step() - with pytest.raises(RuntimeError, match="stale quantized cache"): - module(hidden, routing, sglang_ep_native_local_ids=local_ids) - with pytest.raises(RuntimeError, match="stale quantized cache"): - module.publishable_expert_bytes() - module.refresh_quantized_cache() - module(hidden, routing, sglang_ep_native_local_ids=local_ids) - - -@pytest.mark.gpu -@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires Hopper CUDA") -def test_cuda_straight_through_gradients_match_direct_vjp_with_expert_locality() -> None: - device = _hopper_or_skip() - pytest.importorskip("sglang") - module = _bank(device) - module.load_prequantized(*_checkpoint_bytes(device)) - - rows = 5 - hidden = ( - torch.arange(rows * _HIDDEN, dtype=torch.float32, device=device) - .reshape(rows, _HIDDEN) - .remainder_(23) - .sub_(11) - .div_(64) - .to(torch.bfloat16) - .requires_grad_(True) - ) - local_ids = torch.tensor([[2], [9], [9], [14], [-1]], dtype=torch.int32, device=device) - routing = torch.tensor([[0.75], [0.5], [0.25], [1.0], [1.0]], dtype=torch.float32, device=device).requires_grad_( - True - ) - - output = module(hidden, routing, sglang_ep_native_local_ids=local_ids, routed_scaling_factor=1.5) - grad_output = torch.ones_like(output) - output.backward(grad_output) - - expected = module._straight_through_vjp( - hidden.detach(), - routing.detach(), - local_ids, - grad_output=grad_output, - routed_scaling_factor=1.5, - needs_input_grad=(True, True, True, True), - ) - assert torch.equal(hidden.grad, expected[0].to(torch.bfloat16)) - assert torch.equal(routing.grad, expected[1]) - assert torch.equal(module.gate_up_weight_master.grad, expected[2]) - assert torch.equal(module.down_weight_master.grad, expected[3]) - - routed = {2, 9, 14} - for expert_index in range(_LOCAL_EXPERTS): - nonzero = torch.count_nonzero(module.gate_up_weight_master.grad[expert_index]) - if expert_index in routed: - assert nonzero > 0 - else: - assert nonzero == 0 - - -@pytest.mark.gpu -@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires Hopper CUDA") -def test_cuda_ep_rank_local_banks_partition_values_and_own_master_gradients() -> None: - """Verify EP ownership and normalization. - - Production canonical-EP semantics: every rank's bank sees the GATHERED - token set with foreign slots mapped to the -1 sentinel and routing - weights normalized ONCE globally (sigmoid -> topk -> norm -> scale, - never renormalized over the local expert subset); the combine SUMS the - rank partials. Gate: two EP-local half banks seeded from the same - bytes as a 16-expert reference bank must (a) emit exact zeros for - fully-foreign rows, (b) sum to the reference value, and (c) own - exactly their slice of the master gradients, bitwise. - """ - - device = _hopper_or_skip() - pytest.importorskip("sglang") - torch.manual_seed(29) - - reference = _bank(device) # 16 experts: the single-rank reference - with torch.no_grad(): - reference.gate_up_weight_master.copy_(torch.randn(_LOCAL_EXPERTS, _HIDDEN, 2 * _INTERMEDIATE, device=device)) - reference.down_weight_master.copy_(torch.randn(_LOCAL_EXPERTS, _INTERMEDIATE, _HIDDEN, device=device)) - reference.refresh_quantized_cache() - published = reference.publishable_expert_bytes() - - half = _LOCAL_EXPERTS // 2 - rank_banks: list[Glm52FullParamBlockFP8RoutedExperts] = [] - for rank in (0, 1): - bank = Glm52FullParamBlockFP8RoutedExperts(half, _HIDDEN, _INTERMEDIATE, device=device) - rows = slice(rank * half, (rank + 1) * half) - bank.load_prequantized(published[0][rows], published[1][rows], published[2][rows], published[3][rows]) - bank.assign_global_expert_range(rank * half, _LOCAL_EXPERTS) - rank_banks.append(bank) - # Expert-boundary-preserving quantization: the half banks hold the - # reference's exact bytes, so any value drift below is EP mechanics. - for rank, bank in enumerate(rank_banks): - rows = slice(rank * half, (rank + 1) * half) - assert torch.equal( - bank.gate_up_packed_weight_f32.view(torch.uint8), - reference.gate_up_packed_weight_f32[rows].view(torch.uint8), - ) - - # Gathered token set, topk=2: one row per global expert (sentinel second - # slot with a NONZERO weight that must be ignored), one cross-rank row, - # one expert-multiplicity row. Routing is production-shaped: sigmoid - # scores, normalized over the row's REAL experts once, globally. - rows = _LOCAL_EXPERTS + 2 - cross_row, multiplicity_row = _LOCAL_EXPERTS, _LOCAL_EXPERTS + 1 - hidden_values = ( - torch.arange(rows * _HIDDEN, dtype=torch.float32, device=device) - .reshape(rows, _HIDDEN) - .remainder_(19) - .sub_(9) - .div_(64) - .to(torch.bfloat16) - ) - global_ids = torch.full((rows, 2), -1, dtype=torch.int32, device=device) - global_ids[:_LOCAL_EXPERTS, 0] = torch.arange(_LOCAL_EXPERTS, dtype=torch.int32, device=device) - global_ids[cross_row] = torch.tensor([3, 11], dtype=torch.int32, device=device) - global_ids[multiplicity_row, 0] = 9 - scores = torch.sigmoid(torch.randn(rows, 2, device=device, dtype=torch.float32)) - real = (global_ids >= 0).float() - routing = (scores * real) / ((scores * real).sum(dim=-1, keepdim=True) + 1e-20) - routing = routing + 0.33 * (1.0 - real) # sentinel-slot weights must be dead - routing = routing.contiguous() - - def local_ids_for(rank: int) -> torch.Tensor: - owned = (global_ids >= rank * half) & (global_ids < (rank + 1) * half) - return torch.where(owned, global_ids - rank * half, global_ids.new_full((), -1)).contiguous() - - grad_output = ( - torch.arange(rows * _HIDDEN, dtype=torch.float32, device=device) - .reshape(rows, _HIDDEN) - .remainder_(13) - .sub_(6) - .div_(32) - .to(torch.bfloat16) - ) - - def run(bank: Glm52FullParamBlockFP8RoutedExperts, ids: torch.Tensor): - hidden_leaf = hidden_values.clone().requires_grad_(True) - routing_leaf = routing.clone().requires_grad_(True) - value = bank(hidden_leaf, routing_leaf, sglang_ep_native_local_ids=ids, routed_scaling_factor=1.5) - value.backward(grad_output) - return value.detach(), hidden_leaf.grad, routing_leaf.grad - - reference_value, reference_hidden_grad, reference_routing_grad = run(reference, global_ids) - partials = [run(bank, local_ids_for(rank)) for rank, bank in enumerate(rank_banks)] - - # (a) fully-foreign rows are EXACT zeros (the combine adds partials). - for rank, (value, _, _) in enumerate(partials): - foreign = (local_ids_for(rank) < 0).all(dim=-1) - assert bool(foreign.any()) - assert torch.count_nonzero(value[foreign]) == 0, f"rank {rank} leaked non-zero foreign rows" - - # (b) the combine reproduces the reference value: bitwise wherever a row - # is served by ONE rank; the cross-rank row differs from a single-rank - # reference by exactly the combine's extra BF16 rounding (the reference - # kernel sums both expert contributions before the cast — one rounding; - # the EP combine adds two already-rounded partials — two roundings; the - # production NCCL combine has the same property). The bound below is - # expressed in terms of the rounded operands. - combined = partials[0][0] + partials[1][0] - single_owner = torch.ones(rows, dtype=torch.bool, device=device) - single_owner[cross_row] = False - assert torch.equal(combined[single_owner], reference_value[single_owner]) - cross_diff = (combined[cross_row].float() - reference_value[cross_row].float()).abs() - # Rounding error of the two-partial combine is bounded by the OPERAND - # magnitudes (cancellation makes a result-relative bound wrong). - cross_bound = (partials[0][0][cross_row].float().abs() + partials[1][0][cross_row].float().abs()).clamp( - min=1.0 - ) * 2**-7 - assert bool((cross_diff <= cross_bound).all()), "cross-rank row exceeded one BF16 combine rounding" - - # (c) master-gradient ownership: each rank owns exactly its slice. - assert torch.count_nonzero(reference.gate_up_weight_master.grad) > 0 - for rank, bank in enumerate(rank_banks): - rows_slice = slice(rank * half, (rank + 1) * half) - assert torch.equal(bank.gate_up_weight_master.grad, reference.gate_up_weight_master.grad[rows_slice]) - assert torch.equal(bank.down_weight_master.grad, reference.down_weight_master.grad[rows_slice]) - - # Routing-weight gradients partition per (row, slot): owner slots carry - # the reference gradient, foreign and sentinel slots are exactly zero. - assert torch.equal(partials[0][2] + partials[1][2], reference_routing_grad) - sentinel_slots = global_ids < 0 - assert torch.count_nonzero(reference_routing_grad[sentinel_slots]) == 0 - - # Hidden-state gradients: single-owner rows are bitwise; the cross-rank - # row sums two BF16 partials (two roundings) against the reference's - # single FP32 accumulation (one rounding) — compared at BF16 tolerance. - hidden_grad_sum = partials[0][1].float() + partials[1][1].float() - assert torch.equal(hidden_grad_sum[single_owner].to(torch.bfloat16), reference_hidden_grad[single_owner]) - assert torch.allclose( - hidden_grad_sum[cross_row].to(torch.bfloat16).float(), - reference_hidden_grad[cross_row].float(), - rtol=2e-2, - atol=2e-3, - ) - - -@pytest.mark.gpu -@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires Hopper CUDA") -def test_cuda_production_optimizer_covers_only_masters_and_post_update_equality() -> None: - """Verify optimizer coverage and post-update equality. - - A production-style optimizer built from ``parameters()`` with the - requires_grad filter must cover exactly the two FP32 masters; a real - step must leave the byte caches untouched (bytes move only at refresh), - trip the staleness gate, and after refresh the published bytes must be - the quantization of the UPDATED masters — with the frozen receiver - byte- and value-equal through BOTH transport forms (fused bank payload - and per-expert checkpoint items). - """ - - device = _hopper_or_skip() - pytest.importorskip("sglang") - from xorl.models.transformers.glm5.native_fp8 import Glm52NativeExpertSlotReceiver - - torch.manual_seed(31) - module = _bank(device) - module.load_prequantized(*_checkpoint_bytes(device)) - module.assign_global_expert_range(0, _LOCAL_EXPERTS) - - trainable = [parameter for parameter in module.parameters() if parameter.requires_grad] - assert {id(p) for p in trainable} == {id(module.gate_up_weight_master), id(module.down_weight_master)} - optimizer = torch.optim.AdamW(trainable, lr=0.05) - - cache_before = tuple( - getattr(module, name).detach().clone() - for name in ( - "gate_up_packed_weight_f32", - "gate_up_weight_scale_inv", - "down_packed_weight_f32", - "down_weight_scale_inv", - ) - ) - masters_before = ( - module.gate_up_weight_master.detach().clone(), - module.down_weight_master.detach().clone(), - ) - - hidden, routing, local_ids = _routing_fixture(device) - value = module(hidden, routing, sglang_ep_native_local_ids=local_ids) - value.float().square().sum().backward() - optimizer.step() - - # The step moved ONLY the masters; the consumed bytes are untouched. - for name, before in zip( - ( - "gate_up_packed_weight_f32", - "gate_up_weight_scale_inv", - "down_packed_weight_f32", - "down_weight_scale_inv", - ), - cache_before, - strict=True, - ): - assert torch.equal(getattr(module, name).detach(), before), f"optimizer step mutated {name}" - assert not torch.equal(module.gate_up_weight_master.detach(), masters_before[0]) - assert not torch.equal(module.down_weight_master.detach(), masters_before[1]) - with pytest.raises(RuntimeError, match="stale quantized cache"): - module(hidden, routing, sglang_ep_native_local_ids=local_ids) - - module.refresh_quantized_cache() - published = module.publishable_expert_bytes() - # The refreshed cache is the quantization of the UPDATED masters. - expected_gate_up, expected_gate_up_scale = quantize_expert_masters_to_serving_bytes(module.gate_up_weight_master) - assert torch.equal(published[0].view(torch.uint8), expected_gate_up.view(torch.uint8)) - assert torch.equal(published[1], expected_gate_up_scale) - # ... and it really moved off the step-0 checkpoint bytes. - assert not torch.equal(published[0].view(torch.uint8), _checkpoint_bytes(device)[0].view(torch.uint8)) - - # Post-update equality through BOTH transport forms. - payload = publish_glm52_fullparam_payload( - [ - ("experts", module), - *( - (f"experts_ckpt.{global_id}", publication) - for global_id, publication in module.checkpoint_publications() - ), - ], - weight_version="post-step-1", - weight_step=1, - ) - fused_receiver = Glm52NativeBlockFP8Experts(_LOCAL_EXPERTS, _HIDDEN, _INTERMEDIATE, device=device) - slot_receiver_bank = Glm52NativeBlockFP8Experts(_LOCAL_EXPERTS, _HIDDEN, _INTERMEDIATE, device=device) - - def resolver(target: str, kind: str): - if kind == "block_fp8_expert_bank": - return fused_receiver - assert kind == "block_fp8_expert" - return Glm52NativeExpertSlotReceiver(slot_receiver_bank, int(target.rsplit(".", 1)[1])) - - apply_glm52_fullparam_payload( - payload, - resolver, - expected_inventory=_expected_inventory(payload), - version_guard=Glm52WeightVersionGuard(), - ) - - for receiver in (fused_receiver, slot_receiver_bank): - assert torch.equal( - receiver.gate_up_packed_weight_f32.view(torch.uint8), - module.gate_up_packed_weight_f32.view(torch.uint8), - ) - assert torch.equal( - receiver.down_packed_weight_f32.view(torch.uint8), - module.down_packed_weight_f32.view(torch.uint8), - ) - receiver_value = receiver(hidden, routing, sglang_ep_native_local_ids=local_ids) - trainer_value = module(hidden, routing, sglang_ep_native_local_ids=local_ids) - assert torch.equal(trainer_value.detach(), receiver_value) - - -@pytest.mark.gpu -@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires Hopper CUDA") -def test_cuda_payload_applies_expert_bank_to_frozen_receiver_with_byte_equality() -> None: - device = _hopper_or_skip() - pytest.importorskip("sglang") - torch.manual_seed(17) - module = _bank(device) - with torch.no_grad(): - module.gate_up_weight_master.copy_(torch.randn(_LOCAL_EXPERTS, _HIDDEN, 2 * _INTERMEDIATE, device=device)) - module.down_weight_master.copy_(torch.randn(_LOCAL_EXPERTS, _INTERMEDIATE, _HIDDEN, device=device)) - module.refresh_quantized_cache() - - payload = publish_glm52_fullparam_payload([("experts", module)], weight_version="step-2") - - receiver = Glm52NativeBlockFP8Experts(_LOCAL_EXPERTS, _HIDDEN, _INTERMEDIATE, device=device) - - def resolver(target: str, kind: str): - assert (target, kind) == ("experts", "block_fp8_expert_bank") - return receiver - - apply_glm52_fullparam_payload( - payload, - resolver, - expected_inventory=_expected_inventory(payload), - version_guard=Glm52WeightVersionGuard(), - ) - - hidden, routing, local_ids = _routing_fixture(device) - trainer_value = module(hidden, routing, sglang_ep_native_local_ids=local_ids) - receiver_value = receiver(hidden, routing, sglang_ep_native_local_ids=local_ids) - assert torch.equal(trainer_value.detach(), receiver_value) - - def test_bank_engaged_contract_refuses_understored_rows_before_any_kernel(monkeypatch) -> None: """Stored-rows companion to the frozen bank's admission. diff --git a/tests/models/test_glm52_exact_gate_up_qlora.py b/tests/models/test_glm52_exact_gate_up_qlora.py index 9ca76d42..90c44fdc 100644 --- a/tests/models/test_glm52_exact_gate_up_qlora.py +++ b/tests/models/test_glm52_exact_gate_up_qlora.py @@ -13,7 +13,6 @@ ) from xorl.models.transformers.glm5.exact_qlora import Glm52ExactTP1BlockFP8QLoRALinear from xorl.ops.block_fp8_native import NativeBlockFP8Linear -from xorl.ops.fused_silu_and_mul import exact_fp32_silu_and_mul def _module() -> Glm52ExactTP1FusedGateUpBlockFP8QLoRA: @@ -311,253 +310,6 @@ def _assert_fused_gate_up_contract_fails_closed_before_sglang_import() -> None: assert after == before -@pytest.mark.gpu -@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires Hopper CUDA") -def test_official_fused_gate_up_literal_bytes_graph_metadata_zero_and_gradients() -> None: - pytest.importorskip("sglang") - if torch.cuda.get_device_capability()[0] != 9: - pytest.skip("the qualified exact GLM-5.2 component requires Hopper") - from sglang.kernels.ops.gemm.gate_up_lora_b import gate_up_lora_b_fwd - from sglang.kernels.ops.gemm.sgemm_lora_a import sgemm_lora_a_fwd - from sglang.srt.batch_invariant_ops.bi_silu_and_mul import fp32_silu_and_mul - from sglang.srt.layers.quantization.fp8_utils import triton_w8a8_block_fp8_linear - from sglang.srt.lora.backend.triton_backend import TritonLoRABackend - from sglang.srt.lora.utils import LoRABatchInfo - - device = torch.device("cuda") - rows, in_features, intermediate_size = 17, 6144, 12288 - gate_weight = torch.full( - (intermediate_size, in_features), - 0.25, - dtype=torch.float8_e4m3fn, - device=device, - ) - up_weight = torch.full( - (intermediate_size, in_features), - -0.125, - dtype=torch.float8_e4m3fn, - device=device, - ) - scale_shape = (intermediate_size // 128, in_features // 128) - gate_scales = torch.full(scale_shape, 0.03125, dtype=torch.float32, device=device) - up_scales = torch.full(scale_shape, 0.0625, dtype=torch.float32, device=device) - module = Glm52ExactTP1FusedGateUpBlockFP8QLoRA( - in_features, - intermediate_size, - device=device, - ) - module.load_gate_up_prequantized(gate_weight, gate_scales, up_weight, up_scales) - with torch.no_grad(): - module.gate_proj.lora_A.copy_( - torch.arange(in_features, dtype=torch.float32, device=device) - .remainder_(37) - .sub_(18) - .div_(1024) - .unsqueeze(0) - ) - module.up_proj.lora_A.copy_( - torch.arange(in_features, dtype=torch.float32, device=device) - .remainder_(43) - .sub_(21) - .div_(1536) - .unsqueeze(0) - ) - module.gate_proj.lora_B.copy_( - torch.arange(intermediate_size, dtype=torch.float32, device=device) - .remainder_(47) - .sub_(23) - .div_(2048) - .unsqueeze(1) - ) - module.up_proj.lora_B.copy_( - torch.arange(intermediate_size, dtype=torch.float32, device=device) - .remainder_(53) - .sub_(26) - .div_(1792) - .unsqueeze(1) - ) - input = ( - torch.arange(rows * in_features, dtype=torch.float32, device=device) - .remainder_(127) - .sub_(63) - .div_(64) - .reshape(rows, in_features) - .to(torch.bfloat16) - ) - input.requires_grad_(True) - - # The first module invocation is the cold component cell; the second is - # the warm cell. Build the independent raw-S4 oracle only afterward. - cold_actual = module(input) - warm_actual = module(input) - effective_gate_A = module.gate_proj.lora_A.detach().to(torch.bfloat16).contiguous() - effective_gate_B = module.gate_proj.lora_B.detach().to(torch.bfloat16).contiguous() - effective_up_A = module.up_proj.lora_A.detach().to(torch.bfloat16).contiguous() - effective_up_B = module.up_proj.lora_B.detach().to(torch.bfloat16).contiguous() - stacked_A = torch.cat((effective_gate_A, effective_up_A), dim=0).unsqueeze(0).contiguous() - stacked_B = torch.cat((effective_gate_B, effective_up_B), dim=0).unsqueeze(0).contiguous() - eager_info = LoRABatchInfo( - use_cuda_graph=False, - bs=1, - num_segments=1, - seg_indptr=torch.tensor([0, rows], dtype=torch.int32, device=device), - weight_indices=torch.zeros(1, dtype=torch.int32, device=device), - lora_ranks=torch.ones(1, dtype=torch.int32, device=device), - scalings=torch.ones(1, dtype=torch.float32, device=device), - max_len=rows, - seg_lens=torch.tensor([rows], dtype=torch.int32, device=device), - permutation=None, - expected_tokens=rows, - has_active_lora=True, - ) - direct_base = triton_w8a8_block_fp8_linear( - input.detach(), - module.fp8_weight().contiguous(), - [128, 128], - module.weight_scale_inv.contiguous(), - ) - direct_A = sgemm_lora_a_fwd(input.detach(), stacked_A, eager_info, stack_num=2) - expected = gate_up_lora_b_fwd( - direct_A, - stacked_B, - eager_info, - intermediate_size, - base_output=direct_base.clone(), - ) - assert torch.equal(cold_actual.view(torch.uint8), expected.view(torch.uint8)) - assert torch.equal(warm_actual.view(torch.uint8), cold_actual.view(torch.uint8)) - # Serving's exact mode computes the one-round FP32 SwiGLU - # (SiluAndMul.forward_exact, xorl-sglang f10b907d8); the trainer op must - # match a one-round sampler oracle bitwise. - trainer_activation = exact_fp32_silu_and_mul(cold_actual) - sampler_activation = fp32_silu_and_mul(expected) - assert torch.equal(trainer_activation.view(torch.uint8), sampler_activation.view(torch.uint8)) - - # Build the exact adapter-merged metadata used after S4's production - # decode routing. Sixteen live request rows all select adapter slot zero; - # the backend merges them into the first of eight adapter segments while - # retaining the graph's fixed segment arrays. - graph_slots = 16 - max_loras_per_batch = 8 - graph_backend = TritonLoRABackend(max_loras_per_batch=max_loras_per_batch, device=device) - graph_backend.init_cuda_graph_batch_info(max_bs_in_cuda_graph=graph_slots, num_tokens_per_req=1) - graph_backend.batch_info = graph_backend.cuda_graph_batch_info - graph_backend.batch_info.weight_indices[:graph_slots].zero_() - graph_backend.batch_info.lora_ranks.zero_() - graph_backend.batch_info.lora_ranks[0] = 1 - graph_backend.batch_info.scalings.zero_() - graph_backend.batch_info.scalings[0] = 1.0 - graph_backend.compute_sgemm_routing(use_cuda_graph=True) - graph_info = graph_backend.sgemm_batch_info - assert graph_info is graph_backend.cuda_graph_sgemm_batch_info - assert graph_info.bs == max_loras_per_batch - assert torch.equal( - graph_info.seg_lens, - torch.tensor([graph_slots] + [0] * (max_loras_per_batch - 1), dtype=torch.int32, device=device), - ) - assert torch.equal( - graph_info.seg_indptr, - torch.tensor([0] + [graph_slots] * max_loras_per_batch, dtype=torch.int32, device=device), - ) - assert torch.equal(graph_info.weight_indices, torch.arange(max_loras_per_batch, dtype=torch.int32, device=device)) - assert torch.equal(graph_info.permutation, torch.arange(graph_slots, dtype=torch.int32, device=device)) - assert graph_info.max_len == graph_slots - - graph_input = input.detach()[:graph_slots].contiguous() - graph_base = triton_w8a8_block_fp8_linear( - graph_input, - module.fp8_weight().contiguous(), - [128, 128], - module.weight_scale_inv.contiguous(), - ) - graph_stacked_A = torch.zeros((max_loras_per_batch, *stacked_A.shape[1:]), dtype=stacked_A.dtype, device=device) - graph_stacked_B = torch.zeros((max_loras_per_batch, *stacked_B.shape[1:]), dtype=stacked_B.dtype, device=device) - graph_stacked_A[0].copy_(stacked_A[0]) - graph_stacked_B[0].copy_(stacked_B[0]) - graph_A = sgemm_lora_a_fwd(graph_input, graph_stacked_A, graph_info, stack_num=2) - graph_output = graph_backend.run_gate_up_lora( - graph_input, - graph_stacked_A, - graph_stacked_B, - base_output=graph_base.clone(), - ) - - graph_eager_info = LoRABatchInfo( - use_cuda_graph=False, - bs=1, - num_segments=1, - seg_indptr=torch.tensor([0, graph_slots], dtype=torch.int32, device=device), - weight_indices=torch.zeros(1, dtype=torch.int32, device=device), - lora_ranks=torch.ones(1, dtype=torch.int32, device=device), - scalings=torch.ones(1, dtype=torch.float32, device=device), - max_len=graph_slots, - seg_lens=torch.tensor([graph_slots], dtype=torch.int32, device=device), - permutation=None, - expected_tokens=graph_slots, - has_active_lora=True, - ) - graph_eager_A = sgemm_lora_a_fwd(graph_input, stacked_A, graph_eager_info, stack_num=2) - graph_eager_output = gate_up_lora_b_fwd( - graph_eager_A, - stacked_B, - graph_eager_info, - intermediate_size, - base_output=graph_base.clone(), - ) - assert torch.equal(graph_A.view(torch.uint8), graph_eager_A.view(torch.uint8)) - assert torch.equal(graph_output.view(torch.uint8), graph_eager_output.view(torch.uint8)) - - grad_output = ( - torch.arange(rows * 2 * intermediate_size, dtype=torch.float32, device=device) - .remainder_(61) - .sub_(30) - .div_(31) - .reshape(rows, 2 * intermediate_size) - .to(torch.bfloat16) - ) - base_weight = module._dequantize_base_weight().to(torch.bfloat16) - gate_base_input = input.detach().clone().requires_grad_(True) - up_base_input = input.detach().clone().requires_grad_(True) - gate_base_output = F.linear(gate_base_input, base_weight[:intermediate_size]) - up_base_output = F.linear(up_base_input, base_weight[intermediate_size:]) - gate_grad_output, up_grad_output = grad_output.split(intermediate_size, dim=-1) - torch.autograd.backward((gate_base_output, up_base_output), (gate_grad_output, up_grad_output)) - gate_lora_input = input.detach().float().requires_grad_(True) - up_lora_input = input.detach().float().requires_grad_(True) - reference_factors = tuple( - factor.float().requires_grad_(True) - for factor in (effective_gate_A, effective_gate_B, effective_up_A, effective_up_B) - ) - gate_output = F.linear(F.linear(gate_lora_input, reference_factors[0]), reference_factors[1]) - up_output = F.linear(F.linear(up_lora_input, reference_factors[2]), reference_factors[3]) - torch.cat((gate_output, up_output), dim=-1).backward(grad_output.float()) - expected_gate_dx = gate_base_input.grad.float() + gate_lora_input.grad - expected_up_dx = up_base_input.grad.float() + up_lora_input.grad - expected_dx = expected_gate_dx.to(torch.bfloat16) + expected_up_dx.to(torch.bfloat16) - - cold_actual.backward(grad_output) - - assert torch.equal(input.grad, expected_dx) - for master, reference_factor in zip( - ( - module.gate_proj.lora_A, - module.gate_proj.lora_B, - module.up_proj.lora_A, - module.up_proj.lora_B, - ), - reference_factors, - strict=True, - ): - assert torch.equal(master.grad, reference_factor.grad) - - module.zero_grad(set_to_none=True) - with torch.no_grad(): - for name in module.logical_factor_names: - dict(module.named_parameters())[name].zero_() - zero_output = module(input.detach()) - assert torch.equal(zero_output.view(torch.uint8), direct_base.view(torch.uint8)) - - def test_fused_gate_up_cpu_contract(monkeypatch) -> None: _assert_fused_gate_up_contract_is_one_native_leaf_with_four_logical_fp32_factors() _assert_fused_gate_up_loader_makes_gate_then_up_order_explicit_and_strict() diff --git a/tests/models/test_glm52_exact_lm_head_qlora.py b/tests/models/test_glm52_exact_lm_head_qlora.py index bcc9e274..a79d7544 100644 --- a/tests/models/test_glm52_exact_lm_head_qlora.py +++ b/tests/models/test_glm52_exact_lm_head_qlora.py @@ -23,7 +23,6 @@ _selected_logprob_reference_grad_partitioned, glm52_lm_head_shard, ) -from xorl.ops.bi_families_v2 import exact_temperature_scale_fp32_logits def _component(tp_rank: int = 0, tp_group=None) -> Glm52ExactTP16LmHeadSelectedLogprob: @@ -378,157 +377,3 @@ def _assert_tp_group_validation_rejects_size_order_rank_and_backend(monkeypatch) state["backend"] = "gloo" with pytest.raises(RuntimeError, match="must use NCCL"): component._validate_tp_group() - - -@pytest.mark.gpu -@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires Hopper CUDA") -def test_official_local_shard_literal_v2_bytes_tail_and_surrogate_gradients() -> None: - pytest.importorskip("sglang") - if torch.cuda.get_device_capability()[0] != 9: - pytest.skip("the qualified GLM-5.2 exact LM-head component requires Hopper") - - from sglang.kernels.ops.gemm.sgemm_lora_a import sgemm_lora_a_fwd - from sglang.kernels.ops.gemm.sgemm_lora_b import sgemm_lora_b_fwd - from sglang.srt.batch_invariant_ops import ( - head_v2_full_logits_with_lse, - head_v2_selected_logprob_from_logits, - ) - from sglang.srt.layers.vocab_parallel_embedding import VocabParallelEmbedding - - for rank in range(GLM52_LM_HEAD_TP_SIZE): - expected = glm52_lm_head_shard(rank) - actual = VocabParallelEmbedding._get_indices( - GLM52_LM_HEAD_PADDED_VOCAB_SIZE, - GLM52_LM_HEAD_PADDED_VOCAB_SIZE, - GLM52_LM_HEAD_VOCAB_SIZE, - GLM52_LM_HEAD_VOCAB_SIZE, - rank, - GLM52_LM_HEAD_TP_SIZE, - ) - assert ( - actual.org_vocab_start_index, - actual.org_vocab_end_index, - actual.padded_org_vocab_start_index, - actual.padded_org_vocab_end_index, - actual.num_org_vocab_padding, - ) == ( - expected.vocab_start, - expected.vocab_end, - expected.padded_vocab_start, - expected.padded_vocab_end, - 0, - ) - - device = torch.device("cuda") - component = _component(0).to(device) - rows = 2 - torch.manual_seed(20260807) - hidden = torch.empty((rows, GLM52_LM_HEAD_HIDDEN_SIZE), dtype=torch.bfloat16, device=device).uniform_(-0.125, 0.125) - local_weight = torch.empty( - (GLM52_LM_HEAD_LOCAL_VOCAB_SIZE, GLM52_LM_HEAD_HIDDEN_SIZE), - dtype=torch.bfloat16, - device=device, - ).uniform_(-0.0625, 0.0625) - lora_A = ( - torch.arange(GLM52_LM_HEAD_HIDDEN_SIZE, dtype=torch.float32, device=device) - .sub_(3_071) - .div_(16_384) - .reshape(1, GLM52_LM_HEAD_HIDDEN_SIZE) - .requires_grad_(True) - ) - lora_B = ( - torch.arange(GLM52_LM_HEAD_LOCAL_VOCAB_SIZE, dtype=torch.float32, device=device) - .sub_(4_839) - .div_(32_768) - .reshape(GLM52_LM_HEAD_LOCAL_VOCAB_SIZE, 1) - .requires_grad_(True) - ) - hidden_bytes = hidden.view(torch.uint8).clone() - weight_bytes = local_weight.view(torch.uint8).clone() - A_bytes = lora_A.view(torch.uint8).clone() - B_bytes = lora_B.view(torch.uint8).clone() - effective_A = lora_A.detach().to(torch.bfloat16).contiguous() - effective_B = lora_B.detach().to(torch.bfloat16).contiguous() - - batch_info = lm_head_impl._single_adapter_lm_head_batch_info(device.index, rows) - direct_base, _direct_lse = head_v2_full_logits_with_lse(hidden, local_weight) - direct_a = sgemm_lora_a_fwd(hidden, effective_A.unsqueeze(0), batch_info) - direct_delta = sgemm_lora_b_fwd(direct_a, effective_B.unsqueeze(0), batch_info) - expected_local = sgemm_lora_b_fwd( - direct_a, - effective_B.unsqueeze(0), - batch_info, - base_output=direct_base.clone(), - ) - actual_local = component._exact_local_logits(hidden, local_weight, effective_A, effective_B) - warm_local = component._exact_local_logits(hidden, local_weight, effective_A, effective_B) - - assert direct_base.dtype is torch.float32 - assert direct_a.dtype is torch.bfloat16 - assert direct_delta.dtype is torch.bfloat16 - assert torch.equal(expected_local.view(torch.uint8), (direct_base + direct_delta.float()).view(torch.uint8)) - assert torch.equal(actual_local.view(torch.uint8), expected_local.view(torch.uint8)) - assert torch.equal(warm_local.view(torch.uint8), actual_local.view(torch.uint8)) - assert torch.equal(hidden.view(torch.uint8), hidden_bytes) - assert torch.equal(local_weight.view(torch.uint8), weight_bytes) - assert torch.equal(lora_A.view(torch.uint8), A_bytes) - assert torch.equal(lora_B.view(torch.uint8), B_bytes) - assert torch.equal(effective_A.view(torch.uint8), lora_A.detach().to(torch.bfloat16).view(torch.uint8)) - assert torch.equal(effective_B.view(torch.uint8), lora_B.detach().to(torch.bfloat16).view(torch.uint8)) - - stacked = torch.stack( - [actual_local + torch.tensor(rank / 128, dtype=torch.float32, device=device) for rank in range(16)] - ) - gathered = _rank_order_vocab_from_stacked( - stacked, - expected_world_size=16, - expected_local_vocab_size=GLM52_LM_HEAD_LOCAL_VOCAB_SIZE, - ) - expected_gathered = torch.cat([stacked[rank] for rank in range(16)], dim=-1) - assert torch.equal(gathered.view(torch.uint8), expected_gathered.view(torch.uint8)) - - token_ids = torch.tensor([0, GLM52_LM_HEAD_VOCAB_SIZE - 1], dtype=torch.int64, device=device) - for temperature in ( - None, - torch.ones(2, dtype=torch.float32, device=device), - torch.tensor([0.7, 1.3], dtype=torch.float32, device=device), - ): - actual_logprob = component._selected_logprob_from_gathered( - gathered, - token_ids, - temperature, - ) - score_logits = gathered if temperature is None else exact_temperature_scale_fp32_logits(gathered, temperature) - expected_logprob, _, _ = head_v2_selected_logprob_from_logits( - score_logits, - token_ids, - temperature=None, - ) - assert torch.equal(actual_logprob.view(torch.uint8), expected_logprob.view(torch.uint8)) - - grad_local_logits = ( - torch.arange(rows * GLM52_LM_HEAD_LOCAL_VOCAB_SIZE, dtype=torch.float32, device=device) - .remainder_(127) - .sub_(63) - .div_(64) - .reshape(rows, GLM52_LM_HEAD_LOCAL_VOCAB_SIZE) - ) - grad_hidden, grad_A, grad_B = _local_qlora_surrogate_vjp( - hidden, - local_weight, - effective_A, - effective_B, - grad_local_logits, - needs_input_grad=(True, True, True), - ) - base_hidden = hidden.detach().clone().requires_grad_(True) - F.linear(base_hidden, local_weight).backward(grad_local_logits.to(torch.bfloat16)) - lora_hidden = hidden.float().detach().requires_grad_(True) - reference_A = effective_A.float().detach().requires_grad_(True) - reference_B = effective_B.float().detach().requires_grad_(True) - F.linear(F.linear(lora_hidden, reference_A), reference_B).backward(grad_local_logits) - - assert grad_hidden.dtype is grad_A.dtype is grad_B.dtype is torch.float32 - assert torch.equal(grad_hidden, base_hidden.grad.float() + lora_hidden.grad) - assert torch.equal(grad_A, reference_A.grad) - assert torch.equal(grad_B, reference_B.grad) diff --git a/tests/models/test_glm52_exact_routed_experts_qlora.py b/tests/models/test_glm52_exact_routed_experts_qlora.py index a4bb552b..0408bf32 100644 --- a/tests/models/test_glm52_exact_routed_experts_qlora.py +++ b/tests/models/test_glm52_exact_routed_experts_qlora.py @@ -2,7 +2,6 @@ import pytest import torch -import torch.nn.functional as F from xorl.models.transformers.glm5.exact_routed_experts_qlora import ( GLM52_EXACT_EP16_ROUTED_QLORA_CONTRACT_VERSION, @@ -28,52 +27,6 @@ def _module(owner: int, device: torch.device | str = "cpu") -> Glm52ExactEP16Blo ) -def _load_zero_base(module: Glm52ExactEP16BlockFP8QLoRARoutedExperts) -> None: - device = module.gate_up_packed_weight_f32.device - module.load_prequantized( - torch.zeros( - _LOCAL_EXPERTS, - _HIDDEN, - 2 * _INTERMEDIATE, - dtype=torch.float8_e4m3fn, - device=device, - ), - torch.ones(_LOCAL_EXPERTS, 1, 2, dtype=torch.float32, device=device), - torch.zeros( - _LOCAL_EXPERTS, - _INTERMEDIATE, - _HIDDEN, - dtype=torch.float8_e4m3fn, - device=device, - ), - torch.ones(_LOCAL_EXPERTS, 1, 1, dtype=torch.float32, device=device), - ) - - -def _load_distinguishable_base(module: Glm52ExactEP16BlockFP8QLoRARoutedExperts) -> None: - device = module.gate_up_packed_weight_f32.device - gate_up = torch.empty( - _LOCAL_EXPERTS, - _HIDDEN, - 2 * _INTERMEDIATE, - dtype=torch.float8_e4m3fn, - device=device, - ) - gate_up[..., :_INTERMEDIATE] = 0.015625 - gate_up[..., _INTERMEDIATE:] = 0.03125 - module.load_prequantized( - gate_up, - torch.ones(_LOCAL_EXPERTS, 1, 2, dtype=torch.float32, device=device), - torch.full( - (_LOCAL_EXPERTS, _INTERMEDIATE, _HIDDEN), - 0.015625, - dtype=torch.float8_e4m3fn, - device=device, - ), - torch.ones(_LOCAL_EXPERTS, 1, 1, dtype=torch.float32, device=device), - ) - - def _fill_distinguishable_factors(module: Glm52ExactEP16BlockFP8QLoRARoutedExperts) -> None: owner = module.ep_rank global_ids = torch.arange(_GLOBAL_EXPERTS, dtype=torch.float32, device=module.gate_proj_lora_A.device) @@ -90,74 +43,6 @@ def _bits(tensor: torch.Tensor) -> torch.Tensor: return tensor.contiguous().view(torch.uint16) -def _standalone_hybrid_routed_vjp( - module: Glm52ExactEP16BlockFP8QLoRARoutedExperts, - hidden: torch.Tensor, - routing: torch.Tensor, - local_ids: torch.Tensor, - effective_factors: tuple[torch.Tensor, ...], - grad_output: torch.Tensor, - *, - routed_scaling_factor: float, -) -> tuple[torch.Tensor, ...]: - """Independent staged-QloRA oracle; do not reuse the module surrogate.""" - - from sglang.srt.layers.quantization.fp8_utils import block_quant_dequant - - gate_up_weight = block_quant_dequant( - module.gate_up_proj.transpose(1, 2), - module.gate_up_weight_scale_inv.transpose(1, 2), - [128, 128], - torch.bfloat16, - ) - down_weight = block_quant_dequant( - module.down_proj.transpose(1, 2), - module.down_weight_scale_inv.transpose(1, 2), - [128, 128], - torch.bfloat16, - ) - - with torch.enable_grad(), torch.autocast(device_type=hidden.device.type, enabled=False): - references = [ - hidden.float().detach().requires_grad_(True), - routing.float().detach().requires_grad_(True), - *(factor.float().detach().requires_grad_(True) for factor in effective_factors), - ] - reference_hidden, reference_routing = references[:2] - gate_A, gate_B, up_A, up_B, down_A, down_B = references[2:] - reference_output = reference_hidden * 0.0 - - for local_expert in range(_LOCAL_EXPERTS): - pair_rows, pair_topk = (local_ids == local_expert).nonzero(as_tuple=True) - if pair_rows.numel() == 0: - continue - global_expert = module.expert_offset + local_expert - expert_input = reference_hidden.index_select(0, pair_rows) - base_gate_up = F.linear(expert_input.to(torch.bfloat16), gate_up_weight[local_expert]) - base_gate, base_up = base_gate_up.split(_INTERMEDIATE, dim=-1) - gate_delta = (expert_input @ gate_A[0]) @ gate_B[global_expert] - up_delta = (expert_input @ up_A[0]) @ up_B[global_expert] - gate = (base_gate.float() + gate_delta).to(torch.bfloat16) - up = (base_up.float() + up_delta).to(torch.bfloat16) - activated = F.silu(gate.float()).to(torch.bfloat16) * up - base_down = F.linear(activated, down_weight[local_expert]) - down_delta = (activated.float() @ down_A[global_expert]) @ down_B[0] - down = (base_down.float() + down_delta).to(torch.bfloat16).float() - scores = reference_routing[pair_rows, pair_topk].unsqueeze(1) - reference_output = reference_output.index_add( - 0, - pair_rows, - down * scores * routed_scaling_factor, - ) - - return torch.autograd.grad( - reference_output, - references, - grad_outputs=grad_output.float(), - allow_unused=False, - ) - - def test_routed_bank_topology_remap_and_physical_buffer_policy() -> None: module = _module(7) @@ -294,254 +179,3 @@ def _assert_post_ep_owner_local_factor_banks_produce_same_views() -> None: assert global_buffers.keys() == local_buffers.keys() for name in global_buffers: assert torch.equal(global_buffers[name], local_buffers[name]), name - - -@pytest.mark.gpu -@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires Hopper CUDA") -def test_routed_experts_literal_sampler_and_gradient_policy() -> None: - _assert_all_256_owner_slots_run_literal_sampler_partials_routing_outputs_and_logical_vjps() - _assert_routed_gradient_edge_and_mixed_owner_policy() - - -def _assert_all_256_owner_slots_run_literal_sampler_partials_routing_outputs_and_logical_vjps() -> None: - pytest.importorskip("sglang") - if torch.cuda.get_device_capability()[0] != 9: - pytest.skip("the qualified GLM-5.2 routed component requires Hopper") - device = torch.device("cuda") - global_grid = torch.arange(_GLOBAL_EXPERTS, dtype=torch.int64, device=device).reshape(_EP_SIZE, _LOCAL_EXPERTS) - - for owner in range(_EP_SIZE): - module = _module(owner, device) - _load_distinguishable_base(module) - global_ids = global_grid[owner].reshape(_LOCAL_EXPERTS, 1).contiguous() - hidden = ( - torch.arange(_LOCAL_EXPERTS * _HIDDEN, dtype=torch.float32, device=device) - .reshape(_LOCAL_EXPERTS, _HIDDEN) - .remainder_(17) - .add_(1) - .div_(32) - .to(torch.bfloat16) - .requires_grad_(True) - ) - routing = ((global_ids.float() + 1) / 512).contiguous().requires_grad_(True) - base_output = None - if owner == 0: - with torch.no_grad(): - for name in module.logical_factor_names: - getattr(module, name).zero_() - base_output = module(hidden.detach(), routing.detach(), selected_experts=global_ids).detach() - _fill_distinguishable_factors(module) - output = module(hidden, routing, selected_experts=global_ids) - - assert torch.count_nonzero(output) > 0 - if base_output is not None: - assert not torch.equal(output.detach(), base_output) - - # A positive logical cotangent keeps every intentionally routed slot - # distinguishable; an alternating cotangent can legitimately cancel - # one rank-one bank gradient and would not be a dispatch failure. - grad_output = torch.ones_like(output) - local_ids = module.localize_global_expert_ids(global_ids) - effective = tuple( - getattr(module, name).detach().to(torch.bfloat16).contiguous() for name in module.logical_factor_names - ) - expected_gradients = module._surrogate_vjp( - hidden.detach(), - routing.detach(), - local_ids, - *effective, - grad_output=grad_output, - routed_scaling_factor=1.0, - needs_input_grad=(True, True, False, True, True, True, True, True, True), - ) - assert all(gradient is not None and gradient.dtype is torch.float32 for gradient in expected_gradients) - output.backward(grad_output) - assert hidden.grad is not None and torch.count_nonzero(hidden.grad) > 0 - assert routing.grad is not None and routing.grad.dtype is torch.float32 - assert torch.equal(hidden.grad, expected_gradients[0].to(torch.bfloat16)) - assert torch.equal(routing.grad, expected_gradients[1]) - for name, expected in zip(module.logical_factor_names, expected_gradients[2:], strict=True): - assert torch.equal(getattr(module, name).grad, expected) - for name in ("gate_proj_lora_B", "up_proj_lora_B", "down_proj_lora_A"): - gradient = getattr(module, name).grad - assert gradient is not None - assert gradient.dtype is torch.float32 - counts = torch.count_nonzero(gradient.reshape(_GLOBAL_EXPERTS, -1), dim=1) - owner_start = module.expert_offset - owner_end = owner_start + _LOCAL_EXPERTS - assert torch.all(counts[owner_start:owner_end] > 0) - assert torch.count_nonzero(counts[:owner_start]) == 0 - assert torch.count_nonzero(counts[owner_end:]) == 0 - for name in ("gate_proj_lora_A", "up_proj_lora_A", "down_proj_lora_B"): - gradient = getattr(module, name).grad - assert gradient is not None and gradient.dtype is torch.float32 - assert torch.count_nonzero(gradient) > 0 - - if owner == 0: - half_output = module( - hidden.detach(), - (routing.detach() * 0.5).contiguous(), - selected_experts=global_ids, - ).detach() - torch.testing.assert_close( - half_output.float(), - output.detach().float() * 0.5, - rtol=0, - atol=2**-8, - ) - - -def _assert_routed_gradient_edge_and_mixed_owner_policy() -> None: - pytest.importorskip("sglang") - if torch.cuda.get_device_capability()[0] != 9: - pytest.skip("the qualified GLM-5.2 routed component requires Hopper") - device = torch.device("cuda") - module = _module(5, device) - _load_zero_base(module) - _fill_distinguishable_factors(module) - hidden = torch.full((2, _HIDDEN), 0.25, dtype=torch.bfloat16, device=device, requires_grad=True) - global_ids = torch.tensor([[80], [80]], dtype=torch.int64, device=device) - routing = torch.tensor([[0.75], [0.5]], dtype=torch.float32, device=device, requires_grad=True) - module(hidden, routing, selected_experts=global_ids).float().sum().backward() - - for name in ("gate_proj_lora_B", "up_proj_lora_B", "down_proj_lora_A"): - gradient = getattr(module, name).grad - assert gradient is not None - assert torch.count_nonzero(gradient[80]) > 0 - assert torch.count_nonzero(gradient[:80]) == 0 - assert torch.count_nonzero(gradient[81:]) == 0 - - hidden_backing = torch.zeros((2, 2 * _HIDDEN), dtype=torch.bfloat16, device=device) - hidden_strided = hidden_backing[:, ::2] - assert not hidden_strided.is_contiguous() - with pytest.raises(ValueError, match="non-empty and contiguous"): - module(hidden_strided, routing.detach(), selected_experts=global_ids) - routing_backing = torch.ones((2, 2), dtype=torch.float32, device=device) - routing_strided = routing_backing[:, ::2] - assert not routing_strided.is_contiguous() - with pytest.raises(ValueError, match="route-major"): - module(hidden.detach(), routing_strided, selected_experts=global_ids) - - _assert_all_sentinel_owner_returns_zero_and_zero_gradients() - _assert_topk8_mixed_owner_hybrid_vjps_match_standalone_reference() - - -def _assert_all_sentinel_owner_returns_zero_and_zero_gradients() -> None: - device = torch.device("cuda") - module = _module(5, device) - _load_zero_base(module) - _fill_distinguishable_factors(module) - hidden = torch.full((2, _HIDDEN), 0.25, dtype=torch.bfloat16, device=device, requires_grad=True) - routing = torch.tensor([[0.75], [0.5]], dtype=torch.float32, device=device, requires_grad=True) - unowned_global_ids = torch.tensor([[0], [1]], dtype=torch.int64, device=device) - - output = module(hidden, routing, selected_experts=unowned_global_ids) - assert torch.count_nonzero(output) == 0 - output.float().sum().backward() - - assert hidden.grad is not None and torch.count_nonzero(hidden.grad) == 0 - assert routing.grad is not None and torch.count_nonzero(routing.grad) == 0 - for name in module.logical_factor_names: - gradient = getattr(module, name).grad - assert gradient is not None and gradient.dtype is torch.float32 - assert torch.count_nonzero(gradient) == 0 - - -def _assert_topk8_mixed_owner_hybrid_vjps_match_standalone_reference() -> None: - pytest.importorskip("sglang") - if torch.cuda.get_device_capability()[0] != 9: - pytest.skip("the qualified GLM-5.2 routed component requires Hopper") - device = torch.device("cuda") - module = _module(5, device) - _load_distinguishable_base(module) - with torch.no_grad(): - for offset, name in enumerate(module.logical_factor_names): - parameter = getattr(module, name) - values = torch.arange(parameter.numel(), dtype=torch.float32, device=device).reshape_as(parameter) - parameter.copy_(values.remainder(31).sub(15).div(512).add((offset + 1) / 4096)) - - hidden = ( - torch.arange(4 * _HIDDEN, dtype=torch.float32, device=device) - .reshape(4, _HIDDEN) - .remainder_(23) - .sub_(11) - .div_(16) - .to(torch.bfloat16) - .requires_grad_(True) - ) - global_ids = torch.tensor( - [ - [80, 1, 81, 2, 82, 3, 83, 4], - [84, 85, 86, 87, 88, 89, 90, 91], - [200, 92, 201, 93, 202, 94, 203, 95], - [0, 1, 2, 3, 4, 5, 6, 7], - ], - dtype=torch.int64, - device=device, - ) - owned = (global_ids >= module.expert_offset) & (global_ids < module.expert_offset + _LOCAL_EXPERTS) - expected_local_ids = torch.where(owned, global_ids - module.expert_offset, -1).to(torch.int32).contiguous() - local_ids = module.localize_global_expert_ids(global_ids) - assert global_ids.shape[1] == 8 - assert torch.equal(local_ids, expected_local_ids) - assert torch.equal( - torch.sort(local_ids[local_ids >= 0]).values, - torch.arange(_LOCAL_EXPERTS, dtype=torch.int32, device=device), - ) - assert torch.count_nonzero(local_ids == -1) == 16 - - routing = ( - torch.arange(global_ids.numel(), dtype=torch.float32, device=device) - .reshape_as(global_ids) - .remainder_(13) - .add_(1) - .div_(17) - .contiguous() - .requires_grad_(True) - ) - grad_output = ( - torch.arange(hidden.numel(), dtype=torch.float32, device=device) - .reshape_as(hidden) - .remainder_(19) - .sub_(9) - .div_(16) - .to(torch.bfloat16) - ) - routed_scaling_factor = 2.5 - effective_factors = tuple( - getattr(module, name).detach().to(torch.bfloat16).contiguous() for name in module.logical_factor_names - ) - expected_gradients = _standalone_hybrid_routed_vjp( - module, - hidden.detach(), - routing.detach(), - local_ids, - effective_factors, - grad_output, - routed_scaling_factor=routed_scaling_factor, - ) - - trainables = (hidden, routing, *(getattr(module, name) for name in module.logical_factor_names)) - output = module( - hidden, - routing, - selected_experts=global_ids, - routed_scaling_factor=routed_scaling_factor, - ) - actual_gradients = torch.autograd.grad(output, trainables, grad_outputs=grad_output) - - assert len(actual_gradients) == len(expected_gradients) == 8 - assert torch.equal(actual_gradients[0], expected_gradients[0].to(torch.bfloat16)) - for actual, expected in zip(actual_gradients[1:], expected_gradients[1:], strict=True): - assert actual.dtype is torch.float32 - assert torch.equal(actual, expected) - - assert torch.count_nonzero(actual_gradients[0][3]) == 0 - assert torch.count_nonzero(actual_gradients[1][~owned]) == 0 - assert torch.count_nonzero(actual_gradients[1][owned]) == owned.sum() - for index in (3, 5, 6): - gradient = actual_gradients[index] - per_expert_nonzero = torch.count_nonzero(gradient.reshape(_GLOBAL_EXPERTS, -1), dim=1) - assert torch.all(per_expert_nonzero[module.expert_offset : module.expert_offset + _LOCAL_EXPERTS] > 0) - assert torch.count_nonzero(per_expert_nonzero[: module.expert_offset]) == 0 - assert torch.count_nonzero(per_expert_nonzero[module.expert_offset + _LOCAL_EXPERTS :]) == 0 diff --git a/tests/models/test_glm52_exact_shared_expert_qlora.py b/tests/models/test_glm52_exact_shared_expert_qlora.py index d6ba82f2..67ad933b 100644 --- a/tests/models/test_glm52_exact_shared_expert_qlora.py +++ b/tests/models/test_glm52_exact_shared_expert_qlora.py @@ -4,25 +4,12 @@ import pytest import torch -import torch.nn.functional as F from torch import nn -from xorl.distributed.canonical_moe import CanonicalMoEGraphMetadata from xorl.models.transformers.glm5.exact_shared_expert_qlora import ( GLM52_EXACT_TP16_SHARED_EXPERT_QLORA_CONTRACT_VERSION, Glm52ExactTP16SharedExpertBlockFP8QLoRA, ) -from xorl.ops.fused_silu_and_mul import exact_fp32_silu_and_mul - - -def _canonical_moe_reference(partials: torch.Tensor, metadata: CanonicalMoEGraphMetadata) -> torch.Tensor: - level = [partials[index] for index in range(partials.shape[0])] - while len(level) > 1: - level = [(level[index] + level[index + 1]).to(torch.bfloat16) for index in range(0, len(level), 2)] - result = level[0] - result = result.clone() - result[~metadata.valid_mask] = 0 - return result def _pattern( @@ -56,24 +43,6 @@ def _fill_factors(module: Glm52ExactTP16SharedExpertBlockFP8QLoRA) -> None: module.down_proj.lora_B.copy_(_pattern((6144, 1), modulus=61, center=30, divisor=2560, device=device)) -def _load_base(module: Glm52ExactTP16SharedExpertBlockFP8QLoRA) -> None: - device = module.gate_proj.packed_weight_f32.device - gate_weight = torch.full((2048, 6144), 0.25, dtype=torch.float8_e4m3fn, device=device) - up_weight = torch.full((2048, 6144), -0.125, dtype=torch.float8_e4m3fn, device=device) - down_weight = torch.full((6144, 2048), 0.0625, dtype=torch.float8_e4m3fn, device=device) - gate_scales = torch.full((16, 48), 0.03125, dtype=torch.float32, device=device) - up_scales = torch.full((16, 48), 0.0625, dtype=torch.float32, device=device) - down_scales = torch.full((48, 16), 0.046875, dtype=torch.float32, device=device) - module.load_prequantized( - gate_weight, - gate_scales, - up_weight, - up_scales, - down_weight, - down_scales, - ) - - def test_shared_expert_construction_and_runtime_admission_policy() -> None: for kwargs, message in ( ({"hidden_size": 4096}, "hidden_size=6144"), @@ -282,267 +251,3 @@ def _assert_shared_expert_native_base_views_use_output_rows_and_input_columns() assert torch.equal(actual.gate_up_scales[1], up_scales[ordinal]) assert torch.equal(actual.down_weight.view(torch.uint8), down_weight[:, start:end].contiguous().view(torch.uint8)) assert torch.equal(actual.down_scales[:, 0], down_scales[:, ordinal]) - - -def _manual_local_vjp( - module: Glm52ExactTP16SharedExpertBlockFP8QLoRA, - input: torch.Tensor, - exact_gate_up: torch.Tensor, - exact_activated: torch.Tensor, - grad_output: torch.Tensor, - ordinal: int, -) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: - start, end = ordinal * 128, (ordinal + 1) * 128 - effective_gate_A = module.gate_proj.lora_A.detach().to(torch.bfloat16) - effective_gate_B = module.gate_proj.lora_B.detach().to(torch.bfloat16) - effective_up_A = module.up_proj.lora_A.detach().to(torch.bfloat16) - effective_up_B = module.up_proj.lora_B.detach().to(torch.bfloat16) - effective_down_A = module.down_proj.lora_A.detach().to(torch.bfloat16) - effective_down_B = module.down_proj.lora_B.detach().to(torch.bfloat16) - - def projection_vjp( - projection, - projection_input, - factor_A, - factor_B, - projection_grad, - *, - output_range=None, - input_range=None, - A_input_range=None, - B_output_range=None, - ): - with torch.enable_grad(), torch.autocast(device_type="cuda", enabled=False): - base_input = projection_input.detach().requires_grad_(True) - base_weight = module._dequantized_partition_weight( - projection, - output_range=output_range, - input_range=input_range, - ).to(base_input.dtype) - base_output = F.linear(base_input, base_weight) - (base_input_grad,) = torch.autograd.grad( - base_output, - base_input, - grad_outputs=projection_grad.to(base_output.dtype), - ) - - lora_input = projection_input.float().detach().requires_grad_(True) - reference_A = factor_A.float().detach().requires_grad_(True) - reference_B = factor_B.float().detach().requires_grad_(True) - physical_A = reference_A if A_input_range is None else reference_A[:, A_input_range[0] : A_input_range[1]] - physical_B = reference_B if B_output_range is None else reference_B[B_output_range[0] : B_output_range[1]] - lora_output = F.linear(F.linear(lora_input, physical_A), physical_B) - lora_input_grad, factor_A_grad, factor_B_grad = torch.autograd.grad( - lora_output, - (lora_input, reference_A, reference_B), - grad_outputs=projection_grad.float(), - ) - return base_input_grad.float() + lora_input_grad, factor_A_grad, factor_B_grad - - down_input_grad, down_A_grad, down_B_grad = projection_vjp( - module.down_proj, - exact_activated, - effective_down_A, - effective_down_B, - grad_output, - input_range=(start, end), - A_input_range=(start, end), - ) - with torch.enable_grad(), torch.autocast(device_type="cuda", enabled=False): - gate_up_input = exact_gate_up.detach().requires_grad_(True) - # Mirror the module's VJP reference: differentiate the one-round FP32 - # SwiGLU program the forward now emits. - activation = exact_fp32_silu_and_mul(gate_up_input) - (gate_up_grad,) = torch.autograd.grad( - activation, - gate_up_input, - grad_outputs=down_input_grad.to(activation.dtype), - ) - gate_grad, up_grad = gate_up_grad.split(128, dim=-1) - gate_input_grad, gate_A_grad, gate_B_grad = projection_vjp( - module.gate_proj, - input, - effective_gate_A, - effective_gate_B, - gate_grad, - output_range=(start, end), - B_output_range=(start, end), - ) - up_input_grad, up_A_grad, up_B_grad = projection_vjp( - module.up_proj, - input, - effective_up_A, - effective_up_B, - up_grad, - output_range=(start, end), - B_output_range=(start, end), - ) - return gate_input_grad.to(input.dtype) + up_input_grad.to(input.dtype), { - "gate_proj.lora_A": gate_A_grad, - "gate_proj.lora_B": gate_B_grad, - "up_proj.lora_A": up_A_grad, - "up_proj.lora_B": up_B_grad, - "down_proj.lora_A": down_A_grad, - "down_proj.lora_B": down_B_grad, - } - - -@pytest.mark.gpu -@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires Hopper CUDA") -def test_official_shared_expert_actual_operands_fold_and_surrogate_vjp() -> None: - pytest.importorskip("sglang") - if torch.cuda.get_device_capability(0)[0] != 9: - pytest.skip("the qualified exact GLM-5.2 shared-expert component requires Hopper") - from sglang.kernels.ops.gemm.gate_up_lora_b import gate_up_lora_b_fwd - from sglang.kernels.ops.gemm.sgemm_lora_a import sgemm_lora_a_fwd - from sglang.kernels.ops.gemm.sgemm_lora_b import sgemm_lora_b_fwd - from sglang.srt.batch_invariant_ops.bi_silu_and_mul import fp32_silu_and_mul - from sglang.srt.distributed.canonical_moe import ( - CanonicalRowSlots, - ) - from sglang.srt.distributed.canonical_moe import ( - canonical_moe_reference as sampler_canonical_moe_reference, - ) - from sglang.srt.layers.quantization.fp8_utils import triton_w8a8_block_fp8_linear - from sglang.srt.lora.utils import LoRABatchInfo - - torch.cuda.set_device(0) - device = torch.device("cuda:0") - module = Glm52ExactTP16SharedExpertBlockFP8QLoRA(device=device) - _load_base(module) - _fill_factors(module) - rows, ordinal = 3, 11 - input = _pattern((rows, 6144), modulus=127, center=63, divisor=64, device=device).to(torch.bfloat16) - effective = tuple( - factor.detach().to(torch.bfloat16).contiguous() - for factor in ( - module.gate_proj.lora_A, - module.gate_proj.lora_B, - module.up_proj.lora_A, - module.up_proj.lora_B, - module.down_proj.lora_A, - module.down_proj.lora_B, - ) - ) - actual_witness = module._exact_forward_value(input, *effective, contributor_ordinal=ordinal) - cold_output = module(input, contributor_ordinal=ordinal) - warm_output = module(input, contributor_ordinal=ordinal) - - factors = module._physical_factor_views_from_effective(*effective, ordinal) - base = module._physical_base_views(ordinal) - batch_info = LoRABatchInfo( - use_cuda_graph=False, - bs=1, - num_segments=1, - seg_indptr=torch.tensor([0, rows], dtype=torch.int32, device=device), - weight_indices=torch.zeros(1, dtype=torch.int32, device=device), - lora_ranks=torch.ones(1, dtype=torch.int32, device=device), - scalings=torch.ones(1, dtype=torch.float32, device=device), - max_len=rows, - seg_lens=torch.tensor([rows], dtype=torch.int32, device=device), - permutation=None, - expected_tokens=rows, - has_active_lora=True, - ) - raw_gate_up_base = triton_w8a8_block_fp8_linear( - input, - base.gate_up_weight, - [128, 128], - base.gate_up_scales, - ) - raw_gate_up_A = sgemm_lora_a_fwd(input, factors.gate_up_A, batch_info, stack_num=2) - raw_gate_up = gate_up_lora_b_fwd( - raw_gate_up_A, - factors.gate_up_B, - batch_info, - 128, - base_output=raw_gate_up_base.clone(), - ) - # Serving's exact mode resolves SiluAndMul.forward_exact to the one-round - # FP32 SwiGLU (xorl-sglang f10b907d8); the raw oracle uses serving's op. - raw_activated = fp32_silu_and_mul(raw_gate_up) - raw_down_base = triton_w8a8_block_fp8_linear( - raw_activated, - base.down_weight, - [128, 128], - base.down_scales, - ) - raw_down_A = sgemm_lora_a_fwd(raw_activated, factors.down_A, batch_info) - raw_output = sgemm_lora_b_fwd( - raw_down_A, - factors.down_B, - batch_info, - base_output=raw_down_base.clone(), - ) - - byte_pairs = { - "gate_up_base": (actual_witness.gate_up_base, raw_gate_up_base), - "gate_up_A": (actual_witness.gate_up_A_output, raw_gate_up_A), - "gate_up_post_add": (actual_witness.gate_up, raw_gate_up), - "activation": (actual_witness.activated, raw_activated), - "down_base": (actual_witness.down_base, raw_down_base), - "down_A": (actual_witness.down_A_output, raw_down_A), - "local_partial": (actual_witness.output, raw_output), - "cold": (cold_output, raw_output), - "warm": (warm_output, raw_output), - } - for name, (actual, expected) in byte_pairs.items(): - assert torch.equal(actual.view(torch.uint8), expected.view(torch.uint8)), name - - # One shared logical state emits sixteen physical partials. The existing - # public canonical owner—not a new shared-expert reduction—performs the - # adjacent-pairwise fold in sampler contributor order. - with torch.no_grad(): - fold_input = input[:1].contiguous() - partials = torch.stack( - [module(fold_input, contributor_ordinal=rank) for rank in range(16)], - dim=0, - ) - metadata = CanonicalMoEGraphMetadata.build( - torch.tensor([0], dtype=torch.int64, device=device), - torch.tensor([0], dtype=torch.int64, device=device), - capacity=1, - ) - canonical = _canonical_moe_reference(partials, metadata) - sampler_slots = CanonicalRowSlots.from_positions( - torch.tensor([0], dtype=torch.int64, device=device), - capacity=1, - ) - sampler_canonical = sampler_canonical_moe_reference(partials, sampler_slots) - assert torch.equal(canonical.view(torch.uint8), sampler_canonical.view(torch.uint8)) - - # Validate one physical producer's custom VJP against an independent, - # staged QLoRA reference using the same effective BF16 factor bytes. - grad_input = input.detach().clone().requires_grad_(True) - grad_witness = module._exact_forward_value(grad_input.detach(), *effective, contributor_ordinal=ordinal) - grad_output = _pattern( - (rows, 6144), - modulus=67, - center=33, - divisor=71, - device=device, - ).to(torch.bfloat16) - expected_input_grad, expected_factor_grads = _manual_local_vjp( - module, - grad_input.detach(), - grad_witness.gate_up, - grad_witness.activated, - grad_output, - ordinal, - ) - module(grad_input, contributor_ordinal=ordinal).backward(grad_output) - - assert torch.equal(grad_input.grad, expected_input_grad.to(torch.bfloat16)) - parameters = dict(module.named_parameters()) - for name, expected in expected_factor_grads.items(): - actual = parameters[name].grad - assert actual is not None, name - assert actual.dtype is torch.float32, name - assert torch.equal(actual, expected), name - start, end = ordinal * 128, (ordinal + 1) * 128 - assert not torch.count_nonzero(module.gate_proj.lora_B.grad[:start]) - assert not torch.count_nonzero(module.gate_proj.lora_B.grad[end:]) - assert not torch.count_nonzero(module.up_proj.lora_B.grad[:start]) - assert not torch.count_nonzero(module.up_proj.lora_B.grad[end:]) - assert not torch.count_nonzero(module.down_proj.lora_A.grad[:, :start]) - assert not torch.count_nonzero(module.down_proj.lora_A.grad[:, end:]) diff --git a/tests/models/test_glm52_fullparam_frozen_trunk_backward.py b/tests/models/test_glm52_fullparam_frozen_trunk_backward.py index 60a7eace..38228b19 100644 --- a/tests/models/test_glm52_fullparam_frozen_trunk_backward.py +++ b/tests/models/test_glm52_fullparam_frozen_trunk_backward.py @@ -18,8 +18,6 @@ from __future__ import annotations -import logging - import pytest import torch @@ -31,7 +29,6 @@ glm52_fullparam_routing_weights_with_grad, ) from xorl.models.transformers.glm5.native_fp8 import ( - GLM52_NATIVE_EXPERTS_FROZEN_DGRAD_CONTRACT_VERSION, Glm52NativeBlockFP8Experts, ) @@ -139,89 +136,6 @@ def test_cuda_frozen_bank_value_bytes_identical_with_and_without_grad_engagement assert torch.count_nonzero(scoring) > 0 -@pytest.mark.gpu -@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires Hopper CUDA") -def test_cuda_frozen_bank_activation_grads_match_trainable_bank_and_mutate_nothing(caplog) -> None: - device = _hopper_or_skip() - pytest.importorskip("sglang") - checkpoint = _checkpoint_bytes(device) - - Glm52NativeBlockFP8Experts._frozen_dgrad_engagement_logged = False - frozen = Glm52NativeBlockFP8Experts(_LOCAL_EXPERTS, _HIDDEN, _INTERMEDIATE, device=device) - frozen.load_prequantized(*checkpoint) - with caplog.at_level(logging.INFO, logger="xorl.models.transformers.glm5.native_fp8"): - frozen.enable_frozen_activation_dgrad() - frozen.enable_frozen_activation_dgrad() # idempotent - engagement = [record for record in caplog.records if "frozen-bank activation dgrad engaged" in record.message] - assert len(engagement) == 1 - assert GLM52_NATIVE_EXPERTS_FROZEN_DGRAD_CONTRACT_VERSION in engagement[0].message - - packed_before = { - name: getattr(frozen, name).detach().view(torch.uint8).clone() - for name in ( - "gate_up_packed_weight_f32", - "gate_up_weight_scale_inv", - "down_packed_weight_f32", - "down_weight_scale_inv", - ) - } - - hidden, routing, local_ids = _grad_fixture(device) - frozen_hidden = hidden.clone().requires_grad_(True) - frozen_routing = routing.clone().requires_grad_(True) - output = frozen( - frozen_hidden, - frozen_routing, - sglang_ep_native_local_ids=local_ids, - routed_scaling_factor=1.5, - ) - grad_output = torch.ones_like(output) - output.backward(grad_output) - assert frozen_hidden.grad is not None and frozen_routing.grad is not None - assert bool(frozen_hidden.grad.abs().sum() > 0) and bool(frozen_routing.grad.abs().sum() > 0) - # The sentinel row's hidden gradient is exactly zero (no expert touched it). - assert torch.count_nonzero(frozen_hidden.grad[4]) == 0 - - # Direct-vjp wiring identity (the autograd boundary passes exactly the - # engaged tensors through). - direct_hidden, direct_routing = frozen._frozen_activation_vjp( - hidden, - routing, - local_ids, - grad_output=grad_output, - routed_scaling_factor=1.5, - needs_input_grad=(True, True), - ) - assert torch.equal(frozen_hidden.grad, direct_hidden) - assert torch.equal(frozen_routing.grad, direct_routing) - - # Cross-implementation: the QUALIFIED trainable bank on identical bytes - # produces bitwise-identical hidden/routing gradients (same checked - # program, same bytes) — the frozen path is that treatment minus wgrad. - trainable = Glm52FullParamBlockFP8RoutedExperts(_LOCAL_EXPERTS, _HIDDEN, _INTERMEDIATE, device=device) - trainable.load_prequantized(*checkpoint) - trainable_hidden = hidden.clone().requires_grad_(True) - trainable_routing = routing.clone().requires_grad_(True) - trainable( - trainable_hidden, - trainable_routing, - sglang_ep_native_local_ids=local_ids, - routed_scaling_factor=1.5, - ).backward(grad_output) - assert torch.equal(frozen_hidden.grad, trainable_hidden.grad) - assert torch.equal(frozen_routing.grad, trainable_routing.grad) - # ... and the trainable bank did produce master grads where the frozen - # bank, by construction, has no master to grade. - assert trainable.gate_up_weight_master.grad is not None - - # Frozen means frozen: no parameter gradients, no byte movement. - for name, parameter in frozen.named_parameters(): - assert parameter.grad is None, f"frozen bank parameter {name} received a gradient" - assert not parameter.requires_grad - for name, before in packed_before.items(): - assert torch.equal(getattr(frozen, name).detach().view(torch.uint8), before), name - - # --------------------------------------------------------------------------- # Routing-weight surrogate # --------------------------------------------------------------------------- diff --git a/tests/models/test_glm52_fullparam_reduced_backward_gate.py b/tests/models/test_glm52_fullparam_reduced_backward_gate.py deleted file mode 100644 index a97f921d..00000000 --- a/tests/models/test_glm52_fullparam_reduced_backward_gate.py +++ /dev/null @@ -1,217 +0,0 @@ -"""Reduced-depth full-model backward gate for GLM-5.2 full-param training. - -Component, EP, and FSDP2 tests do not by themselves prove that gradients -traverse every frozen module between trainable surfaces. This test runs a -real full-depth forward and backward on one GPU at reduced GLM-5.2 geometry -with a scoped trainable set. - -Trainable scope mirrored from production (scoped admission): dense MLP -composites (layer 0), one in-scope routed-expert bank (layer 1), ALL -routers (layers 1, 2). Frozen trunk on the gradient path: attention -projections (q_a/q_b/kv_a/o as ``NativeBlockFP8Linear``; kv_b through the -absorbed dequant-einsum), shared experts (``NativeBlockFP8Linear`` with -block-aligned ranges), the out-of-scope frozen expert bank (layer 2), all -norms, the frozen BF16 LM head. - -Asserts, on one loss backward: -1. every admitted FP32 master receives a finite, nonzero gradient — the - layer-0 dense masters specifically prove the gradient traversed layers - 1-2's ENTIRE frozen trunk; -2. no frozen parameter receives any gradient; -3. no frozen parameter byte and no quantized-cache byte changes — frozen - means frozen: the dgrad mechanism must never write masters or caches. - -Reduced-scale seam (documented fidelity boundary): the canonical MoE -dispatch requires the EP16/CP16 topology -(``_canonical_ep_forward``), and the generic eager MoE path cannot call -the native banks at all, so this gate substitutes a SINGLE-CONTRIBUTOR -dispatch that calls the model's real ``_canonical_routed_local_partial`` -and ``_canonical_shared_local_partial`` seams — every bank / shared-expert -/ projection class boundary (where the backward mechanisms live) is the -production one. The 16-rank ordered combine itself stays gated by -tests/distributed/test_glm52_fullparam_ep16_combine.py; attention uses the -differentiable torch sparse-MLA reference (the flashmla envelope is -official-geometry-only; its trainable backward is gated separately). -""" - -from __future__ import annotations - -import pytest -import torch -import torch.nn.functional as F - -from tests.models.test_glm52_fullparam_admission import ( - _hopper_or_skip, - _seed_native_bytes, - _tiny_config, -) - - -def _gate_config(): - """3-layer reduced geometry: dense + in-scope sparse + out-of-scope sparse.""" - - config = _tiny_config() - config.num_hidden_layers = 3 - config.mlp_layer_types = ["dense", "sparse", "sparse"] - config.indexer_types = ["full", "shared", "shared"] - # Keep the routed scaling out of the reduced seam: the generic trainer - # route pre-multiplies routing weights while the canonical bank call - # applies the factor inside the kernel; 1.0 makes both identity so the - # reduced path cannot double-apply it. - config.routed_scaling_factor = 1.0 - # Production uses sparse MLA. flashmla's envelope is official-geometry-only; - # the torch reference is the differentiable reduced-scale stand-in. - config._sparse_mla_enabled = True - config._sparse_mla_backend = "torch" - quantization = dict(config.quantization_config) - exclusions = [entry for entry in quantization["modules_to_not_convert"] if not entry.startswith("model.layers.")] - exclusions.extend(f"model.layers.{layer}.self_attn.indexers_proj" for layer in range(3)) - quantization["modules_to_not_convert"] = exclusions - config.quantization_config = quantization - return config - - -def _single_contributor_experts_with_shared( - self, - hidden_states: torch.Tensor, - routing_weights: torch.Tensor, - selected_experts: torch.Tensor, - absolute_positions: torch.Tensor | None = None, -): - """EP1 projection of the canonical dispatch through the REAL partial seams.""" - - del absolute_positions - batch_size, seq_len, hidden_dim = hidden_states.shape - flat = hidden_states.reshape(-1, hidden_dim) - rows = flat.shape[0] - routing = routing_weights.reshape(rows, -1).float().contiguous() - # Single contributor owns the full bank: global ids ARE the local slots. - local_ids = selected_experts.reshape(rows, -1).to(torch.int32).contiguous() - routed = self._canonical_routed_local_partial(flat, routing, selected_experts, local_ids) - shared = self._canonical_shared_local_partial(flat, contributor_ordinal=0, contributor_count=1) - return (routed + shared).to(torch.bfloat16).reshape(batch_size, seq_len, hidden_dim) - - -def _byte_snapshot(model: torch.nn.Module) -> dict[str, torch.Tensor]: - """Clone every frozen parameter and every buffer (byte caches included).""" - - snapshot: dict[str, torch.Tensor] = {} - for name, parameter in model.named_parameters(): - if not parameter.requires_grad: - snapshot[f"param:{name}"] = parameter.detach().clone() - for name, buffer in model.named_buffers(): - snapshot[f"buffer:{name}"] = buffer.detach().clone() - return snapshot - - -def _run_full_depth_backward_gate(monkeypatch: pytest.MonkeyPatch, *, exact_program: bool) -> None: - device = _hopper_or_skip() - pytest.importorskip("sglang") - from xorl.models.transformers.glm5.exact_fullparam_admission import ( - install_glm52_fullparam_components, - ) - from xorl.models.transformers.glm5.modeling_glm5 import Glm5ForCausalLM, Glm5MoEBlock - - config = _gate_config() - if exact_program: - # The production resolver's full-param flag selects the exact forward — the - # structural BI router contract, canonical serving grouped top-k in - # route(), Class-B RoPE, and the exact indexer selector — at reduced - # geometry. The fused serving-norm mode rides along. - config._glm52_fullparam_training = True - import xorl.models.layers.normalization as normalization - - # RMSNorm modules capture the mode at __init__; monkeypatch restores - # the module global at teardown. - monkeypatch.setattr(normalization, "_RMSNORM_MODE", "sglang_fused") - torch.manual_seed(9173) - model = Glm5ForCausalLM(config).to(torch.bfloat16).to(device) - if exact_program: - # The exact indexer SELECTOR's kernels (flashinfer layernorm, fused - # BF16 projection) are official-geometry-bound and reject reduced - # shapes; the frozen selector runs entirely under no_grad, so - # it is outside the backward under test. Pin it to the reduced-scale - # legacy path — a documented fidelity boundary of this gate, like - # flashmla and the distributed combine. - from xorl.models.transformers.glm5.modeling_glm5 import Glm5Attention - - for module in model.modules(): - if isinstance(module, Glm5Attention) and module.indexer is not None: - monkeypatch.setattr(module.indexer, "selector_version", "legacy_torch_or_tilelang") - _seed_native_bytes(model) - report = install_glm52_fullparam_components( - model, config, trainable_expert_layers=(1,), _skip_geometry_validation=True - ) - assert report.dense_mlp_layers == (0,) - assert report.routed_expert_layers == (1,) - assert report.router_layers == (1, 2) - - monkeypatch.setattr( - Glm5MoEBlock, - "forward_experts_with_shared", - _single_contributor_experts_with_shared, - ) - - model.train() - trainable = {name: parameter for name, parameter in model.named_parameters() if parameter.requires_grad} - assert set(trainable) == { - "model.layers.0.mlp.gate_up_proj.weight_master", - "model.layers.0.mlp.down_proj.weight_master", - "model.layers.1.mlp.experts.gate_up_weight_master", - "model.layers.1.mlp.experts.down_weight_master", - "model.layers.1.mlp.gate.full_param.weight_master", - "model.layers.2.mlp.gate.full_param.weight_master", - } - frozen = {name: parameter for name, parameter in model.named_parameters() if not parameter.requires_grad} - assert frozen, "reduced model lost its frozen trunk" - - before = _byte_snapshot(model) - - input_ids = torch.randint(0, config.vocab_size, (1, 32), device=device) - outputs = model(input_ids=input_ids, index_share_mode="training_with_backward") - logits = model.lm_head(outputs.last_hidden_state).float() - loss = F.cross_entropy( - logits[:, :-1].reshape(-1, config.vocab_size), - input_ids[:, 1:].reshape(-1), - ) - assert loss.requires_grad, "loss is detached from the trainable masters" - loss.backward() - - # 1. Gradient reaches EVERY admitted master through the full frozen depth. - for name, parameter in trainable.items(): - assert parameter.grad is not None, f"no gradient reached {name}" - assert bool(torch.isfinite(parameter.grad).all()), f"non-finite gradient on {name}" - assert bool(parameter.grad.abs().sum() > 0), f"gradient on {name} is exactly zero" - - # 2. No gradient lands on any frozen parameter. - for name, parameter in frozen.items(): - assert parameter.grad is None, f"frozen parameter {name} received a gradient" - assert not parameter.requires_grad, f"frozen parameter {name} was unfrozen" - - # 3. Frozen bytes and cache bytes are bit-identical after the backward. - after = _byte_snapshot(model) - assert set(after) == set(before) - for name, tensor in before.items(): - assert torch.equal(after[name], tensor), f"bytes changed during forward/backward: {name}" - - -@pytest.mark.gpu -@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires Hopper CUDA") -def test_reduced_full_depth_backward_reaches_all_masters_and_touches_no_frozen_bytes( - monkeypatch: pytest.MonkeyPatch, -) -> None: - _run_full_depth_backward_gate(monkeypatch, exact_program=False) - - -@pytest.mark.gpu -@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires Hopper CUDA") -def test_reduced_full_depth_backward_under_the_exact_forward_program( - monkeypatch: pytest.MonkeyPatch, -) -> None: - """Run the same gate with the production program selection engaged: the - full-param flag (structural BI router contract + canonical serving - grouped top-k + Class-B RoPE + exact indexer selector) and the fused - serving norm mode. FlashMLA and the distributed combine remain outside - this single-GPU projection.""" - - _run_full_depth_backward_gate(monkeypatch, exact_program=True) diff --git a/tests/models/test_qwen3_5_trunk_wrap.py b/tests/models/test_qwen3_5_trunk_wrap.py index bf101e8c..5ace6a77 100644 --- a/tests/models/test_qwen3_5_trunk_wrap.py +++ b/tests/models/test_qwen3_5_trunk_wrap.py @@ -11,7 +11,6 @@ router gate (contracted separately by the exact model program) and lm_head/embed. """ -import pytest import torch from xorl.lora.modules.linear import LoraLinear @@ -26,9 +25,6 @@ ) -requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") - - def _hybrid_config(**overrides) -> Qwen3_5MoeConfig: """Two layers: layer 0 linear-attention + dense MLP, layer 1 full-attention + sparse MoE (shared expert included).""" @@ -137,24 +133,3 @@ def _is_wrapped(module): assert not _is_wrapped(model.model.embed_tokens) finally: set_trunk_linear_contract(False) - - -@requires_cuda -@pytest.mark.gpu -def test_qwen3_5_full_attn_forward_runs_under_trunk_wrap(): - """Wrapped full-attention + shared-expert + dense projections must run the - bf16 contract GEMM end-to-end (the runtime guard raises on any non-bf16 - operand).""" - torch.manual_seed(1) - config = _hybrid_config(layer_types=["full_attention", "full_attention"], _moe_implementation="eager") - model = Qwen3_5MoeForCausalLM(config).to(device="cuda", dtype=torch.bfloat16).eval() - try: - wrapped = wrap_trunk_linears_batch_invariant(model) - assert wrapped["shared_expert_gate"] == 1 - input_ids = torch.randint(0, config.vocab_size, (1, 8), device="cuda") - with torch.no_grad(): - out = model(input_ids=input_ids) - assert out.last_hidden_state.dtype == torch.bfloat16 - assert torch.isfinite(out.last_hidden_state.float()).all() - finally: - set_trunk_linear_contract(False) diff --git a/tests/models/test_qwen3_moe_fused_lora.py b/tests/models/test_qwen3_moe_fused_lora.py index c2922745..bdf8cbaa 100644 --- a/tests/models/test_qwen3_moe_fused_lora.py +++ b/tests/models/test_qwen3_moe_fused_lora.py @@ -354,93 +354,6 @@ def _make_pair(self, ref_backend, test_backend, device): _copy_block_weights(ref, test) return ref, test - def test_zero_lora_and_nonzero_lora(self): - """With lora_B=0, LoRA output must equal base; nonzero LoRA must change output.""" - device = "cuda" - for backend in ["eager", "triton", "native"]: - # --- zero LoRA matches base --- - base_block = MoEBlock( - hidden_size=self.HIDDEN_DIM, - num_experts=self.NUM_EXPERTS, - top_k=2, - intermediate_size=self.INTERMEDIATE, - moe_implementation=backend, - ) - torch.manual_seed(42) - nn.init.xavier_normal_(base_block.experts.gate_proj.data) - nn.init.xavier_normal_(base_block.experts.up_proj.data) - nn.init.xavier_normal_(base_block.experts.down_proj.data) - nn.init.xavier_normal_(base_block.gate.weight.data) - base_block = base_block.to(device).to(self.DTYPE) - - lora_block = MoEBlock( - hidden_size=self.HIDDEN_DIM, - num_experts=self.NUM_EXPERTS, - top_k=2, - intermediate_size=self.INTERMEDIATE, - moe_implementation=backend, - ) - torch.manual_seed(42) - nn.init.xavier_normal_(lora_block.experts.gate_proj.data) - nn.init.xavier_normal_(lora_block.experts.up_proj.data) - nn.init.xavier_normal_(lora_block.experts.down_proj.data) - nn.init.xavier_normal_(lora_block.gate.weight.data) - lora_block = lora_block.to(device).to(self.DTYPE) - lora_block.inject_lora(r=self.R, lora_alpha=self.LORA_ALPHA) - - torch.manual_seed(999) - hidden = torch.randn(2, 8, self.HIDDEN_DIM, device=device, dtype=self.DTYPE) - base_out, _ = base_block(hidden) - lora_out, _ = lora_block(hidden) - torch.testing.assert_close( - lora_out, - base_out, - atol=1e-3, - rtol=1e-2, - msg=f"[{backend}] Zero-LoRA output should match base model", - ) - - # --- nonzero LoRA changes output --- - block = _make_lora_block( - backend, - self.NUM_EXPERTS, - self.HIDDEN_DIM, - self.INTERMEDIATE, - self.R, - self.LORA_ALPHA, - device, - self.DTYPE, - ) - with torch.no_grad(): - for proj in ["gate_proj", "up_proj", "down_proj"]: - lora_B = getattr(block.experts, f"{proj}_lora_B") - nn.init.xavier_normal_(lora_B) - - base_block2 = ( - MoEBlock( - hidden_size=self.HIDDEN_DIM, - num_experts=self.NUM_EXPERTS, - top_k=2, - intermediate_size=self.INTERMEDIATE, - moe_implementation=backend, - ) - .to(device) - .to(self.DTYPE) - ) - with torch.no_grad(): - base_block2.gate.weight.copy_(block.gate.weight) - base_block2.experts.gate_proj.copy_(block.experts.gate_proj) - base_block2.experts.up_proj.copy_(block.experts.up_proj) - base_block2.experts.down_proj.copy_(block.experts.down_proj) - - torch.manual_seed(999) - hidden2 = torch.randn(2, 8, self.HIDDEN_DIM, device=device, dtype=self.DTYPE) - base_out2, _ = base_block2(hidden2) - lora_out2, _ = block(hidden2) - - diff = (lora_out2 - base_out2).abs().max().item() - assert diff > 1e-3, f"[{backend}] Non-zero LoRA should change the output, but max diff={diff}" - def test_cross_backend_output_and_gradients(self): """Cross-backend outputs and LoRA gradients should match.""" for ref_backend, test_backend in [("eager", "native"), ("eager", "triton"), ("triton", "native")]: diff --git a/tests/ops/test_bi_fused_lm_head.py b/tests/ops/test_bi_fused_lm_head.py index c0d27de5..7310adbf 100644 --- a/tests/ops/test_bi_fused_lm_head.py +++ b/tests/ops/test_bi_fused_lm_head.py @@ -1,7 +1,6 @@ import pytest import torch -from xorl.ops.loss.bi_fused_lm_head import bi_fused_per_token_ce from xorl.ops.loss.causallm_loss import causallm_loss_function @@ -161,52 +160,6 @@ def test_bi_fused_per_row_unit_temperature_preserves_forward_bytes(): assert torch.equal(scalar.per_token_logprobs, per_row.per_token_logprobs) -@requires_cuda -@pytest.mark.gpu -@pytest.mark.parametrize("family", ["v1", "v2"]) -def test_bi_fused_temperature_matches_serving_materialize_then_score(family): - pytest.importorskip("sglang") - from sglang.srt.batch_invariant_ops import ( - bi_lm_head_selected_logprob_from_logits as serving_v1_score, - ) - from sglang.srt.batch_invariant_ops import ( - exact_temperature_scale_fp32_logits as serving_scale, - ) - from sglang.srt.batch_invariant_ops import ( - head_v2_selected_logprob_from_logits as serving_v2_score, - ) - - from xorl.ops import bi_families_v2 - from xorl.ops.batch_invariant_ops import bi_lm_head_full_logits - from xorl.ops.bi_families_v2 import head_v2_full_logits_with_lse - - torch.manual_seed(53) - hidden = torch.randn((4, 128), dtype=torch.bfloat16, device="cuda") - weight = torch.randn((512, 128), dtype=torch.bfloat16, device="cuda") - labels = torch.tensor([1, 127, 255, 511], dtype=torch.int64, device="cuda") - temperature = torch.tensor([0.7, 1.0, 1.3, 0.9], dtype=torch.float32, device="cuda") - try: - if family == "v1": - bi_families_v2._select_qwen35_families_v1() - logits = bi_lm_head_full_logits(hidden, weight) - score = serving_v1_score - else: - bi_families_v2._select_glm52_families_v2() - logits, _ = head_v2_full_logits_with_lse(hidden, weight) - score = serving_v2_score - - actual = bi_fused_per_token_ce(hidden, weight, labels, temperature=temperature) - transformed = serving_scale(logits, temperature) - expected_logprob, _, _ = score(transformed, labels, temperature=None) - assert torch.equal(actual.view(torch.uint8), (-expected_logprob).view(torch.uint8)) - - scalar_unit = bi_fused_per_token_ce(hidden, weight, labels, temperature=1.0) - row_unit = bi_fused_per_token_ce(hidden, weight, labels, temperature=torch.ones_like(temperature)) - assert torch.equal(scalar_unit.view(torch.uint8), row_unit.view(torch.uint8)) - finally: - bi_families_v2._select_nonexact_families() - - @requires_cuda @pytest.mark.gpu def test_bi_kernel_unit_temperature_is_exact_identity(): diff --git a/tests/ops/test_eager_vs_native_moe.py b/tests/ops/test_eager_vs_native_moe.py deleted file mode 100644 index 887d3948..00000000 --- a/tests/ops/test_eager_vs_native_moe.py +++ /dev/null @@ -1,193 +0,0 @@ -"""Correctness comparison between eager and native MoE backends. - -Tests forward-pass output agreement, backward-pass gradient agreement, -determinism, and scaling behavior across varying expert counts, hidden dims, -top-k values, and batch/sequence sizes. -""" - -import pytest -import torch -import torch.nn as nn - - -DEVICE = "cuda" -DTYPE = torch.bfloat16 - - -# --------------------------------------------------------------------------- -# Helpers (lazy-import to avoid torchvision env crash at module level) -# --------------------------------------------------------------------------- - - -def _import_moe(): - """Import MoE layers lazily so CPU-only collection stays lightweight.""" - from xorl.models.layers.moe.experts import MoEExperts # noqa: PLC0415 - from xorl.models.layers.moe.moe_block import MoEBlock # noqa: PLC0415 - - return MoEBlock, MoEExperts - - -def _make_pair(num_experts, hidden_dim, intermediate, top_k, seed=42): - """Create eager + native MoEBlocks with identical weights.""" - MoEBlock, _ = _import_moe() - - torch.manual_seed(seed) - eager = MoEBlock(hidden_dim, num_experts, top_k, intermediate, moe_implementation="eager") - nn.init.xavier_normal_(eager.experts.gate_proj.data) - nn.init.xavier_normal_(eager.experts.up_proj.data) - nn.init.xavier_normal_(eager.experts.down_proj.data) - nn.init.xavier_normal_(eager.gate.weight.data) - eager = eager.to(DEVICE, DTYPE) - - native = MoEBlock(hidden_dim, num_experts, top_k, intermediate, moe_implementation="native") - native = native.to(DEVICE, DTYPE) - with torch.no_grad(): - native.gate.weight.copy_(eager.gate.weight) - native.experts.gate_proj.copy_(eager.experts.gate_proj) - native.experts.up_proj.copy_(eager.experts.up_proj) - native.experts.down_proj.copy_(eager.experts.down_proj) - return eager, native - - -# --------------------------------------------------------------------------- -# Test 1: Forward + backward agreement across all configs -# --------------------------------------------------------------------------- - -PARITY_CONFIGS = [ - # (num_experts, hidden_dim, intermediate, top_k, batch, seq) - (4, 64, 128, 2, 2, 8), - (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 seq -] - - -def _assert_forward_and_backward_agreement(ne, hd, inter, topk, bs, seq): - eager_block, native_block = _make_pair(ne, hd, inter, topk) - config = f"E={ne}, H={hd}, I={inter}, top_k={topk}, batch={bs}, seq={seq}" - - # --- Forward agreement --- - torch.manual_seed(999) - x = torch.randn(bs, seq, hd, device=DEVICE, dtype=DTYPE) - with torch.no_grad(): - eager_out, eager_logits = eager_block(x) - native_out, native_logits = native_block(x) - - torch.testing.assert_close(eager_logits, native_logits, atol=0, rtol=0) - max_diff = (eager_out - native_out).abs().max().item() - torch.testing.assert_close( - native_out, - eager_out, - atol=0.05, - rtol=0.02, - msg=f"Forward mismatch ({config}): max_diff={max_diff:.6f}", - ) - - # --- Backward agreement --- - torch.manual_seed(999) - x_eager = torch.randn(bs, seq, hd, device=DEVICE, dtype=DTYPE, requires_grad=True) - x_native = x_eager.detach().clone().requires_grad_(True) - - eager_out2, _ = eager_block(x_eager) - eager_out2.sum().backward() - native_out2, _ = native_block(x_native) - native_out2.sum().backward() - - atol, rtol = 0.05, 0.05 - torch.testing.assert_close( - x_native.grad, - x_eager.grad, - atol=atol, - rtol=rtol, - msg=f"Input gradient mismatch ({config})", - ) - for name in ["gate_proj", "up_proj", "down_proj"]: - eager_grad = getattr(eager_block.experts, name).grad - native_grad = getattr(native_block.experts, name).grad - assert eager_grad is not None, f"eager {name} grad is None ({config})" - assert native_grad is not None, f"native {name} grad is None ({config})" - torch.testing.assert_close( - native_grad, - eager_grad, - atol=atol, - rtol=rtol, - msg=f"{name} gradient mismatch ({config})", - ) - torch.testing.assert_close( - native_block.gate.weight.grad, - eager_block.gate.weight.grad, - atol=atol, - rtol=rtol, - msg=f"Gate weight gradient mismatch ({config})", - ) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -def test_forward_and_backward_agreement(): - """Eager and native agree at routing and token-count boundaries.""" - for config in PARITY_CONFIGS: - _assert_forward_and_backward_agreement(*config) - - _assert_determinism_and_edge_cases() - - -# --------------------------------------------------------------------------- -# Test 2: Determinism + edge case (all tokens same expert) -# --------------------------------------------------------------------------- - - -def _assert_determinism_and_edge_cases(): - """Determinism: same input produces identical output. Edge case: all tokens to same expert.""" - MoEBlock, MoEExperts = _import_moe() - - for backend in ["eager", "native"]: - # --- Determinism --- - torch.manual_seed(42) - block = MoEBlock(64, 4, 2, 128, moe_implementation=backend) - nn.init.xavier_normal_(block.experts.gate_proj.data) - nn.init.xavier_normal_(block.experts.up_proj.data) - nn.init.xavier_normal_(block.experts.down_proj.data) - nn.init.xavier_normal_(block.gate.weight.data) - block = block.to(DEVICE, DTYPE) - - torch.manual_seed(999) - x = torch.randn(2, 8, 64, device=DEVICE, dtype=DTYPE) - with torch.no_grad(): - out1, _ = block(x) - out2, _ = block(x) - assert torch.equal(out1, out2), f"{backend} is not deterministic" - - # --- All tokens to same expert --- - from xorl.models.layers.moe.backend.native import native_expert_forward # noqa: PLC0415 - - ne, hd, inter = 4, 64, 128 - torch.manual_seed(42) - experts = MoEExperts(ne, hd, inter, moe_implementation="eager").to(DEVICE, DTYPE) - nn.init.xavier_normal_(experts.gate_proj.data) - nn.init.xavier_normal_(experts.up_proj.data) - nn.init.xavier_normal_(experts.down_proj.data) - - num_tokens, top_k = 16, 2 - x_edge = torch.randn(num_tokens, hd, device=DEVICE, dtype=DTYPE) - routing_weights = torch.ones(num_tokens, top_k, device=DEVICE, dtype=DTYPE) / top_k - selected_experts = torch.zeros(num_tokens, top_k, device=DEVICE, dtype=torch.long) - - with torch.no_grad(): - native_out = native_expert_forward( - x_edge, - routing_weights, - selected_experts, - experts.gate_proj, - experts.up_proj, - experts.down_proj, - num_experts=ne, - ) - eager_out = experts(x_edge, expert_idx=0) - - torch.testing.assert_close( - native_out, - eager_out, - atol=0.01, - rtol=0.01, - msg="Same-expert output mismatch", - ) diff --git a/tests/ops/test_expert_kernel_layout_invariance.py b/tests/ops/test_expert_kernel_layout_invariance.py index 34b61316..3e6df51d 100644 --- a/tests/ops/test_expert_kernel_layout_invariance.py +++ b/tests/ops/test_expert_kernel_layout_invariance.py @@ -105,24 +105,6 @@ def _active_slot_bits(per_slot: torch.Tensor, ids: torch.Tensor) -> torch.Tensor return _bits(per_slot[active]) -def test_subset_composition_invariance(fixture): - """Dedup-receive subset vs the all-tokens batch: same rows, same bytes.""" - hidden, ids, weights, run = fixture["hidden"], fixture["ids"], fixture["weights"], fixture["run"] - full_slots = run(hidden, ids, weights, no_combine=True) - full_partial = run(hidden, ids, weights, no_combine=False) - - subset = (ids >= 0).any(dim=1) - sub_slots = run(hidden[subset], ids[subset], weights[subset], no_combine=True) - sub_partial = run(hidden[subset], ids[subset], weights[subset], no_combine=False) - - assert torch.equal(_active_slot_bits(sub_slots, ids[subset]), _active_slot_bits(full_slots[subset], ids[subset])), ( - "per-slot expert bytes depend on batch composition (subset vs all-tokens)" - ) - assert torch.equal(_bits(sub_partial), _bits(full_partial[subset])), ( - "per-token local partial depends on batch composition (subset vs all-tokens)" - ) - - def test_filler_row_invariance(fixture): """Interleaved all(-1) filler rows (capacity padding) must not perturb real rows.""" hidden, ids, weights, run = fixture["hidden"], fixture["ids"], fixture["weights"], fixture["run"] @@ -175,75 +157,6 @@ def test_m_sweep_row_stability(fixture): ) -def test_local_combine_is_a_one_way_boundary(fixture): - """The fused kernel's internal local combine - cannot be reconstructed from its own per-slot (no_combine) outputs. - - Source reading (fused_moe.py): with no_combine=False the routing weight - is applied inside the down-GEMM epilogue in FP32 (mul_routed_weight) - and each WEIGHTED slot is rounded to BF16 once (intermediate_cache3); - moe_sum_reduce then reduces those rows. The no_combine=True output is - instead the UNWEIGHTED row rounded to BF16 — the FP32 pre-weight value - is internal-only, so no external composition of the per-slot rows can - reproduce the combined bytes: bf16(w * row_fp32) != any f(bf16(row)). - - Consequence for the contract: the shared all-tokens program's local - combine is a THIRD numerical program (neither canonical_combine_v1 nor - reconstructable from canonical inputs). An exact DeepEP path must pair - no_combine per-slot rows with canonical_combine on BOTH engines (the - serving ep_gather program), not attempt to mimic this internal combine. - - The gate asserts all four external reconstructions FAIL. If one ever - MATCHES, the kernel's rounding boundary moved — reclassify before - trusting any existing byte contract. - """ - hidden, ids, weights, run = fixture["hidden"], fixture["ids"], fixture["weights"], fixture["run"] - per_slot = run(hidden, ids, weights, no_combine=True) - partial = run(hidden, ids, weights, no_combine=False) - - active = ids >= 0 - weighted_fp32 = per_slot.float() * weights.unsqueeze(-1) - weighted_fp32 = torch.where(active.unsqueeze(-1), weighted_fp32, torch.zeros_like(weighted_fp32)) - # Hypotheses 3/4: the kernel applies the routing weight in the down-GEMM - # epilogue (mul_routed_weight) and ROUNDS EACH WEIGHTED SLOT TO BF16 - # before any summation (intermediate_cache3 is bf16). The reduction then - # runs over pre-rounded slot rows. - weighted_bf16 = weighted_fp32.to(torch.bfloat16) - - candidates = { - "fp32_products_fp32_acc": weighted_fp32.sum(dim=1).to(torch.bfloat16), - "bf16_rounded_slots_fp32_acc": weighted_bf16.float().sum(dim=1).to(torch.bfloat16), - } - acc = torch.zeros_like(partial, dtype=torch.bfloat16) - for k in range(_K): - acc = (acc.float() + weighted_fp32[:, k]).to(torch.bfloat16) - candidates["fp32_products_bf16_chain"] = acc - acc = torch.zeros_like(partial, dtype=torch.bfloat16) - for k in range(_K): - acc = (acc.float() + weighted_bf16[:, k].float()).to(torch.bfloat16) - candidates["bf16_rounded_slots_bf16_chain"] = acc - - matches = {name: bool(torch.equal(_bits(value), _bits(partial))) for name, value in candidates.items()} - print(f"fused-kernel local combine external reconstructions (all expected False): {matches}") - assert not any(matches.values()), ( - "ROUNDING-BOUNDARY CHANGE: an external reconstruction of the fused kernel's local " - f"combine now matches ({matches}); the weight-before-round boundary documented in " - "this test and the contract doc has moved — reclassify before trusting byte contracts" - ) - # Sanity that the mismatch is a rounding boundary, not garbage. An - # elementwise band is the wrong instrument here: slot cancellation can - # blow up single-element relative error. Use per-row relative L2, which is cancellation-tolerant and - # still an order stricter than any non-boundary failure mode. - closest = candidates["bf16_rounded_slots_fp32_acc"].float() - row_norm = partial.float().norm(dim=1) - meaningful = row_norm > 1.0 - rel = (closest - partial.float()).norm(dim=1)[meaningful] / row_norm[meaningful] - assert float(rel.max()) < 2e-2, ( - f"the reconstruction is not even numerically close (max row rel-L2 {float(rel.max()):.4f}); " - "something beyond the documented rounding boundary differs" - ) - - def test_weight_application_point(fixture): """Pin where routing weights are applied: per-slot rows must be UNWEIGHTED. diff --git a/tests/ops/test_moe_torch_compile.py b/tests/ops/test_moe_torch_compile.py deleted file mode 100644 index 07d1112b..00000000 --- a/tests/ops/test_moe_torch_compile.py +++ /dev/null @@ -1,207 +0,0 @@ -"""Test torch.compile compatibility for MoE models. - -Tests per-layer compilation (like torchtitan's apply_compile) on: -- MoEBlock alone (native/eager/triton/quack backends) -- Qwen3MoeDecoderLayer -- Full Qwen3MoeForCausalLM forward + backward -- TFLOPS benchmark: compiled vs uncompiled - -Known issues: -- fullgraph=True: graph break from logger.warning_once in get_parallel_state() -- triton/quack backends: custom autograd.Function causes graph breaks but - works with fullgraph=False (torch.compile splits around the opaque kernels). -""" - -import pytest -import torch -import torch.nn as nn - -from xorl.models.layers.rope import RotaryEmbedding -from xorl.models.transformers.qwen3_moe.configuration_qwen3_moe import Qwen3MoeConfig -from xorl.models.transformers.qwen3_moe.modeling_qwen3_moe import ( - Qwen3MoeDecoderLayer, - Qwen3MoeForCausalLM, -) - - -DEVICE = "cuda" -DTYPE = torch.bfloat16 - - -def _tiny_moe_config(**overrides): - """Create a minimal Qwen3MoeConfig for fast testing.""" - - defaults = dict( - 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="native", - max_position_embeddings=128, - pad_token_id=0, - _attn_implementation="sdpa", - ) - defaults.update(overrides) - return Qwen3MoeConfig(**defaults) - - -def _make_position_embeddings(config, seq_len, device, dtype): - """Create position_embeddings (cos, sin) for decoder layer tests.""" - - rotary = RotaryEmbedding(config=config).to(device) - dummy_hidden = torch.randn(1, seq_len, config.hidden_size, device=device, dtype=dtype) - position_ids = torch.arange(seq_len, device=device).unsqueeze(0) - cos, sin = rotary(dummy_hidden, position_ids) - return cos, sin - - -def _make_moe_block(moe_backend, hidden_size=128, num_experts=4, top_k=2, intermediate=128): - """Create an MoEBlock with xavier init for numerical stability.""" - from xorl.models.layers.moe.moe_block import MoEBlock # noqa: PLC0415 - - block = MoEBlock( - hidden_size=hidden_size, - num_experts=num_experts, - top_k=top_k, - intermediate_size=intermediate, - moe_implementation=moe_backend, - ) - nn.init.xavier_normal_(block.experts.gate_proj.data) - nn.init.xavier_normal_(block.experts.up_proj.data) - nn.init.xavier_normal_(block.experts.down_proj.data) - nn.init.xavier_normal_(block.gate.weight.data) - return block.to(DEVICE, DTYPE) - - -def _available_backends(): - """Return list of available MoE backends on this system.""" - from xorl.utils.import_utils import is_fused_moe_available # noqa: PLC0415 - - backends = ["native", "eager"] - if is_fused_moe_available(): - backends.append("triton") - backends.append("quack") - return backends - - -AVAILABLE_BACKENDS = _available_backends() if torch.cuda.is_available() else [] - - -# --------------------------------------------------------------------------- -# Test 1: MoEBlock compile -- aot_eager + inductor + fullgraph + correctness -# --------------------------------------------------------------------------- - - -def _assert_moe_block_compile(moe_backend): - """MoEBlock compile: aot_eager tracing, inductor compile, and correctness.""" - # --- aot_eager tracing (forward + backward) --- - block = _make_moe_block(moe_backend) - compiled_aot = torch.compile(block, fullgraph=False, backend="aot_eager") - - x = torch.randn(2, 8, 128, device=DEVICE, dtype=DTYPE, requires_grad=True) - out, router_logits = compiled_aot(x) - assert out.shape == x.shape - assert router_logits.shape == (16, 4) - out.sum().backward() - assert x.grad is not None - - # --- inductor compile (forward + backward) --- - block2 = _make_moe_block(moe_backend) - compiled_ind = torch.compile(block2, fullgraph=False, backend="inductor") - x2 = torch.randn(2, 8, 128, device=DEVICE, dtype=DTYPE, requires_grad=True) - out2, _ = compiled_ind(x2) - assert out2.shape == x2.shape - out2.sum().backward() - assert x2.grad is not None - - # --- correctness: compiled vs uncompiled match --- - torch.manual_seed(42) - block4 = _make_moe_block(moe_backend) - x4 = torch.randn(2, 8, 128, device=DEVICE, dtype=DTYPE) - with torch.no_grad(): - ref_out, ref_logits = block4(x4) - compiled_block4 = torch.compile(block4, fullgraph=False, backend="aot_eager") - with torch.no_grad(): - comp_out, comp_logits = compiled_block4(x4) - torch.testing.assert_close(ref_out, comp_out, atol=1e-3, rtol=1e-3) - torch.testing.assert_close(ref_logits, comp_logits, atol=0, rtol=0) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -def test_moe_block_decoder_and_full_model_compile_policy(): - for moe_backend in AVAILABLE_BACKENDS: - _assert_moe_block_compile(moe_backend) - _assert_decoder_layer_compile(moe_backend) - - # Lower-level contracts already compile every available MoE backend with - # both compiler backends. Full-model composition only needs each compiler - # path once. - for moe_backend, compile_backend in (("native", "aot_eager"), ("eager", "inductor")): - _assert_full_model_per_layer_compile(moe_backend, compile_backend) - - -# --------------------------------------------------------------------------- -# Test 2: Qwen3MoeDecoderLayer compile (aot_eager + inductor) -# --------------------------------------------------------------------------- - - -def _assert_decoder_layer_compile(moe_backend): - """Decoder layer compile: aot_eager and inductor, forward + backward.""" - - seq_len = 8 - for compile_backend in ["aot_eager", "inductor"]: - config = _tiny_moe_config(_moe_implementation=moe_backend) - layer = Qwen3MoeDecoderLayer(config, layer_idx=0).to(DEVICE, DTYPE) - compiled_layer = torch.compile(layer, fullgraph=False, backend=compile_backend) - - x = torch.randn(2, seq_len, 128, device=DEVICE, dtype=DTYPE, requires_grad=True) - position_ids = torch.arange(seq_len, device=DEVICE).unsqueeze(0).expand(2, -1) - position_embeddings = _make_position_embeddings(config, seq_len, DEVICE, DTYPE) - - outputs = compiled_layer( - hidden_states=x, - position_ids=position_ids, - position_embeddings=position_embeddings, - ) - - hidden_out = outputs[0] - assert hidden_out.shape == x.shape - hidden_out.sum().backward() - assert x.grad is not None - x.grad = None - - -# --------------------------------------------------------------------------- -# Test 3: Full model per-layer compile (torchtitan style) -# --------------------------------------------------------------------------- - - -def _assert_full_model_per_layer_compile(moe_backend, compile_backend): - """Apply torch.compile to each decoder layer, run forward + backward.""" - - config = _tiny_moe_config(_moe_implementation=moe_backend) - model = Qwen3MoeForCausalLM(config).to(DEVICE, DTYPE) - - compiled_count = 0 - for layer_id, mod in model.model.layers.named_children(): - if isinstance(mod, Qwen3MoeDecoderLayer): - compiled_mod = torch.compile(mod, fullgraph=False, backend=compile_backend) - model.model.layers.register_module(layer_id, compiled_mod) - compiled_count += 1 - - input_ids = torch.randint(0, 1000, (2, 16), device=DEVICE) - - output = model(input_ids=input_ids) - assert output.last_hidden_state is not None - - 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 found"