glm 5.2 lora fp8 rollout fix

[kailash109]

Fix LoRA-B buffer sizing for quantized column-parallel layers

Motivation

Serving a LoRA adapter on an FP8-quantized base (--quantization fp8 --lora-paths ...) crashes at engine init when the adapter targets a fused column-parallel module:

ValueError: LoRA B output dim 4096 does not match base partition prefix dim 2048 for 2 slices

I encountered this error when trying to train GLM-5.2 with LoRA on gate_proj/up_proj, TP > 1 with sglang fp8 rollouts (--quantization fp8). The identical configuration works in bf16 with no fp8 rollout.

Root cause

LoRAMemoryPool.get_lora_B_shape sizes the per-rank LoRA-B buffer by dividing get_hidden_dim()'s full output dim by the TP degree, gated on heuristics that probe base-layer attributes to detect non-standard sharding (e.g. replicated dense/shared-expert gate_up_proj under --moe-dense-tp-size 1).

The probe reads input_size // input_size_per_partition. UnquantizedLinearMethod never sets input_size_per_partition, so on bf16 the probe fails and the fallback path computes the correct per-rank dim. Fp8LinearMethod.create_weights does set layer.input_size_per_partition — and for a column-parallel layer input is unsharded, so the ratio is 1. The probe then reports "unsharded", the split is skipped, and B is allocated at the full output dim against a TP-sharded base. Then the error arises when set_lora_info correctly rejects the mismatch (shared_experts.gate_up_proj: B = 4096 vs per-rank output_partition_sizes = [1024, 1024]).

Modifications

Add _column_parallel_out_partition(): Resolve matching base modele per (module_name, layer_idx) and read output_size_per_partition (same quantity validated against in set_lora_info).

In get_lora_B_shape, for non-MoE column-parallel modules, use the probed per-rank dim directly, and fall back to the existing _column_parallel_lora_b_per_rank_dim arithmetic (which retains the qkv_proj KV-head-replication special case) only when no base module can be probed. MoE modules keep the moe_tp_size path unchanged. For bf16, the probed value equals what previous path computed in all preexisting working configs.

Validation

Pre-fix: crash at lora/layers.py set_lora_info on GLM-5.2 5-layer, TP2, FP8 (module mlp.shared_experts.gate_up_proj, Fp8LinearMethod, B=(4096,16) vs partitions [1024,1024]); same crash at TP32 ([64,64], prefix 128).
Post-fix: same configs load, attach, serve, and run RL rollout + LoRA weight-sync end-to-end at 4-GPU and 64-GPU scale; bf16 runs unchanged.

The training scripts are separate and run Miles on top of Modal -- if necessary I can create a repro of serving GLM5.2 Lora via sglang with and without this fix. Thanks!

---

CI States

Latest PR Test (Base): ⚠️ Run #29280640163
Latest PR Test (Extra): ⚠️ Run #29280639826

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!