[content] miles/backends/fsdp_utils/update_weight_utils.py (python)
def _transform_lora_weights_for_sglang(state_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
    """
    Transform LoRA weight names from separate q_proj/v_proj to combined qkv_proj format for SGLang.
    
    SGLang's Qwen3 model uses combined qkv_proj, but training uses separate q_proj and v_proj.
    This function combines q_proj and v_proj LoRA weights into qkv_proj format.
    
    Args:
        state_dict: Dictionary with keys like "model.layers.0.self_attn.q_proj.lora_A"
                   Values may be DTensors (from FSDP) or regular tensors.
        
    Returns:
        Transformed dictionary with keys like "model.layers.0.self_attn.qkv_proj.lora_A"
        All tensors are converted to regular tensors (not DTensors).
    """
    transformed = {}
    q_proj_weights = {}
    v_proj_weights = {}
    
    # Pattern to match layer indices: model.layers.{layer_idx}.self_attn.{proj_type}.lora_{A|B}
    pattern = re.compile(r"model\.layers\.(\d+)\.self_attn\.(q_proj|v_proj)\.(lora_[AB])")
    
    # Group weights by layer and type, converting DTensors to local tensors
    for name, tensor in state_dict.items():
        # Convert DTensor to local tensor if needed
        if isinstance(tensor, DTensor):
            # Redistribute to Replicate and get local tensor
            tensor = tensor.redistribute(
                placements=[Replicate()] * tensor.device_mesh.ndim,
                async_op=False,
            ).to_local()
        
        match = pattern.match(name)
        if match:
            layer_idx, proj_type, lora_type = match.groups()
            key = (layer_idx, lora_type)
            if proj_type == "q_proj":
                q_proj_weights[key] = tensor
            elif proj_type == "v_proj":
                v_proj_weights[key] = tensor
        else:
            # Keep non-q_proj/v_proj weights as-is (e.g., other LoRA weights)
            transformed[name] = tensor
    
    # Combine q_proj and v_proj into qkv_proj
    # For qkv_proj, we need to interleave Q, K, V. Since we don't have K LoRA weights,
    # we'll create zero tensors for K and combine: [Q_weights, zeros_for_K, V_weights]
    for (layer_idx, lora_type) in set(list(q_proj_weights.keys()) + list(v_proj_weights.keys())):
        q_weight = q_proj_weights.get((layer_idx, lora_type))
        v_weight = v_proj_weights.get((layer_idx, lora_type))
        
        if q_weight is None or v_weight is None:
            logger.warning(
                f"Missing LoRA weight for layer {layer_idx}, type {lora_type}: "
                f"q_proj={q_weight is not None}, v_proj={v_weight is not None}. Skipping."
            )
            continue
        
        # Ensure tensors are on the same device and contiguous
        q_weight = q_weight.contiguous()
        v_weight = v_weight.contiguous()
        
        # For qkv_proj, SGLang expects the weights in QKV order
        # Since we don't have K LoRA weights, we create zeros with the same shape as V (K and V typically have same shape)
        # The exact combination depends on the model architecture, but for now we'll concatenate: Q, zeros(K), V
        # Note: This is a simplified approach - the actual qkv_proj structure may be more complex
        k_zeros = torch.zeros_like(v_weight)
        
        # Combine: Q, K (zeros), V
        # The order and shape depend on how SGLang structures qkv_proj
        # For now, we'll try a simple concatenation along the first dimension
        qkv_weight = torch.cat([q_weight, k_zeros, v_weight], dim=0)
        
        new_name = f"model.layers.{layer_idx}.self_attn.qkv_proj.{lora_type}"
        transformed[new_name] = qkv_weight
        logger.debug(
            f"Transformed LoRA weight: layer {layer_idx}, {lora_type}: "
            f"q_proj shape {q_weight.shape}, v_proj shape {v_weight.shape} -> "
            f"qkv_proj shape {qkv_weight.shape}"
        )
    
    return transformed