sample/jetengine_ext/utils/loader.py

import os
from glob import glob
import torch
from torch import nn
from safetensors import safe_open
import torch.distributed as dist


def default_weight_loader(param: nn.Parameter, loaded_weight: torch.Tensor):
    param.data.copy_(loaded_weight)

def _register_empty_parameter(module, name: str):
    empty = nn.Parameter(torch.empty(0, device='meta'), requires_grad=False)
    module.register_parameter(name, empty)

def _prepare_fused_tensors(model: nn.Module, device: torch.device | str = "cuda"):
    for module in model.modules():
        if not (hasattr(module, "experts") and module.experts):
            continue

        exp0 = module.experts[0]
        if hasattr(exp0, "gate_up_proj"):
            # gate_up_proj weight shape (2 * intermediate_local, hidden_size)
            shape = (len(module.experts),) + exp0.gate_up_proj.weight.shape
            module.register_buffer("_w1",
                                   torch.empty(shape,
                                               dtype=exp0.gate_up_proj.weight.dtype,
                                               device=device),
                                   persistent=False)

        if hasattr(exp0, "down_proj"):
            # down_proj weight shape (hidden_size, intermediate_local)
            shape = (len(module.experts),) + exp0.down_proj.weight.shape
            module.register_buffer("_w2",
                                   torch.empty(shape,
                                               dtype=exp0.down_proj.weight.dtype,
                                               device=device),
                                   persistent=False)

        # Strip per-expert parameters to save VRAM (weights now live in fused buffers)
        for expert in module.experts:
            if hasattr(expert, "gate_up_proj"):
                _register_empty_parameter(expert.gate_up_proj, "weight")
            if hasattr(expert, "down_proj"):
                _register_empty_parameter(expert.down_proj, "weight")

def _is_moe_expert_weight(weight_name: str) -> bool:
    """Check if weight belongs to an MoE expert."""
    return 'experts.' in weight_name and ('gate_up_proj' in weight_name or 'down_proj' in weight_name)

def _load_expert_weight_to_fused(model: nn.Module, weight_name: str, weight_tensor: torch.Tensor, shard_id=None):
    """Load expert weight directly into the appropriate fused tensor with tensor parallel support.

    Tensor Parallel Rules:
      gate_up_proj (ColumnParallel): shard on dim0 (output dimension)
      gate_proj/up_proj shards (when merged loading) each cover half of dim0; we fill into fused _w1 slices
      down_proj (RowParallel): shard on dim1 (input dimension)
    """
    parts = weight_name.split('.')
    layer_path = []
    expert_idx = None
    proj_type = None

    for i, part in enumerate(parts):
        if part == 'experts':
            expert_idx = int(parts[i + 1])
            proj_type = parts[i + 2]  # gate_up_proj or down_proj
            layer_path = parts[:i]
            break

    if expert_idx is None:
        return

    # Resolve module
    moe_module = model
    for attr in layer_path:
        moe_module = getattr(moe_module, attr)

    tp_size = dist.get_world_size() if dist.is_available() and dist.is_initialized() else 1
    tp_rank = dist.get_rank() if dist.is_available() and dist.is_initialized() else 0

    if proj_type == 'gate_up_proj' and hasattr(moe_module, '_w1'):
        fused_tensor = moe_module._w1  # (E, 2*I_local, H)
        local_out = fused_tensor.shape[1]
        # Two cases:
        # 1. Loading merged tensor (gate+up) => shard_id is None, weight_tensor shape (2*I_global, H)
        # 2. Loading individual shard (gate or up) => shard_id in {0,1}, weight_tensor shape (I_global, H)
        if shard_id is None:
            # merged load
            if weight_tensor.shape[0] == local_out:
                # Already sharded
                fused_tensor[expert_idx].copy_(weight_tensor)
            else:
                assert weight_tensor.shape[0] % tp_size == 0 and (weight_tensor.shape[0] // tp_size) == local_out, \
                    f"Unexpected gate_up_proj merged shape {weight_tensor.shape} vs local {fused_tensor.shape} tp={tp_size}"
                local_weight = weight_tensor.narrow(0, tp_rank * local_out, local_out)
                fused_tensor[expert_idx].copy_(local_weight)
        else:
            # individual gate or up
            half_local = local_out // 2
            if weight_tensor.shape[0] == half_local:
                # Already sharded
                start_idx = shard_id * half_local
                fused_tensor[expert_idx, start_idx:start_idx + half_local].copy_(weight_tensor)
            else:
                global_half = weight_tensor.shape[0]
                assert global_half % tp_size == 0 and global_half // tp_size == half_local, \
                    f"Unexpected gate/up proj shard shape {weight_tensor.shape} expected per-rank {half_local}"
                local_weight = weight_tensor.narrow(0, tp_rank * half_local, half_local)
                start_idx = shard_id * half_local
                fused_tensor[expert_idx, start_idx:start_idx + half_local].copy_(local_weight)
    elif proj_type == 'down_proj' and hasattr(moe_module, '_w2'):
        fused_tensor = moe_module._w2  # (E, H, I_local)
        local_in = fused_tensor.shape[2]
        if weight_tensor.shape[1] == local_in:
            fused_tensor[expert_idx].copy_(weight_tensor)
        else:
            assert weight_tensor.shape[1] % tp_size == 0 and (weight_tensor.shape[1] // tp_size) == local_in, \
                f"Unexpected down_proj shape {weight_tensor.shape} vs local {fused_tensor.shape} tp={tp_size}"
            local_weight = weight_tensor.narrow(1, tp_rank * local_in, local_in)
            fused_tensor[expert_idx].copy_(local_weight)

def load_model(model: nn.Module, path: str):
    _prepare_fused_tensors(model)

    packed_modules_mapping = getattr(model, "packed_modules_mapping", {})
    for file in glob(os.path.join(path, "*.safetensors")):
        with safe_open(file, "pt", "cpu") as f:
            for weight_name in f.keys():
                for k in packed_modules_mapping:
                    if k in weight_name:
                        v, shard_id = packed_modules_mapping[k]
                        param_name = weight_name.replace(k, v)
                        if _is_moe_expert_weight(param_name):
                            _load_expert_weight_to_fused(model, param_name, f.get_tensor(weight_name), shard_id)
                        else:
                            param = model.get_parameter(param_name)
                            weight_loader = getattr(param, "weight_loader")
                            weight_loader(param, f.get_tensor(weight_name), shard_id)
                        break
                else:
                    if _is_moe_expert_weight(weight_name):
                        _load_expert_weight_to_fused(model, weight_name, f.get_tensor(weight_name))
                    else:
                        param = model.get_parameter(weight_name)
                        weight_loader = getattr(param, "weight_loader", default_weight_loader)
                        weight_loader(param, f.get_tensor(weight_name))


def load_from_hf_model(target_model: nn.Module, hf_model: nn.Module):
    """
    Load weights from a HuggingFace transformer model to the custom model.
    
    Args:
        target_model: model for inference engine
        hf_model: model from HuggingFace modeling
    """
    import torch
    import torch.nn as nn

    # Get device from HF model
    device = next(hf_model.parameters()).device

    # Track which parameters have been loaded
    loaded_params = set()
    missing_params = []

    # Build a new state dict with materialized tensors
    new_state_dict = {}

    # Process each parameter in the target model
    for name, param in target_model.named_parameters():
        with torch.no_grad():
            # Handle fused qkv_proj
            if 'qkv_proj.weight' in name:
                base_name = name.replace('qkv_proj.weight', '')
                q_name = base_name + 'q_proj.weight'
                k_name = base_name + 'k_proj.weight'
                v_name = base_name + 'v_proj.weight'

                try:
                    q_weight = hf_model.get_parameter(q_name)
                    k_weight = hf_model.get_parameter(k_name)
                    v_weight = hf_model.get_parameter(v_name)

                    # Concatenate Q, K, V weights along dim 0 (output dimension)
                    qkv_weight = torch.cat(
                        [q_weight, k_weight, v_weight], dim=0)

                    # Convert to target dtype and appropriate device
                    if param.is_meta:
                        # Materialize on the HF model's device
                        new_state_dict[name] = qkv_weight.to(
                            device=device, dtype=param.dtype)
                    else:
                        # Keep on parameter's current device
                        new_state_dict[name] = qkv_weight.to(
                            device=param.device, dtype=param.dtype)

                    loaded_params.add(name)
                except AttributeError:
                    missing_params.append(name)
                    # Only keep existing non-meta parameters
                    if not param.is_meta:
                        new_state_dict[name] = param.data

            # Handle fused gate_up_proj (non-MoE)
            elif 'gate_up_proj.weight' in name and 'experts.' not in name:
                base_name = name.replace('gate_up_proj.weight', '')
                gate_name = base_name + 'gate_proj.weight'
                up_name = base_name + 'up_proj.weight'

                try:
                    gate_weight = hf_model.get_parameter(gate_name)
                    up_weight = hf_model.get_parameter(up_name)

                    # Concatenate gate and up weights along dim 0 (output dimension)
                    gate_up_weight = torch.cat([gate_weight, up_weight], dim=0)

                    # Convert to target dtype and appropriate device
                    if param.is_meta:
                        # Materialize on the HF model's device
                        new_state_dict[name] = gate_up_weight.to(
                            device=device, dtype=param.dtype)
                    else:
                        # Keep on parameter's current device
                        new_state_dict[name] = gate_up_weight.to(
                            device=param.device, dtype=param.dtype)

                    loaded_params.add(name)
                except AttributeError:
                    missing_params.append(name)
                    # Only keep existing non-meta parameters
                    if not param.is_meta:
                        new_state_dict[name] = param.data

            # Handle regular parameters (direct mapping)
            else:
                try:
                    hf_param = hf_model.get_parameter(name)

                    # Convert to target dtype and appropriate device
                    if param.is_meta:
                        # Materialize on the HF model's device
                        new_state_dict[name] = hf_param.to(
                            device=device, dtype=param.dtype)
                    else:
                        # Keep on parameter's current device
                        new_state_dict[name] = hf_param.to(
                            device=param.device, dtype=param.dtype)

                    loaded_params.add(name)
                except AttributeError:
                    # Try without model prefix if not found
                    if name.startswith('model.'):
                        try:
                            hf_param = hf_model.get_parameter(name[6:])

                            # Convert to target dtype and appropriate device
                            if param.is_meta:
                                # Materialize on the HF model's device
                                new_state_dict[name] = hf_param.to(
                                    device=device, dtype=param.dtype)
                            else:
                                # Keep on parameter's current device
                                new_state_dict[name] = hf_param.to(
                                    device=param.device, dtype=param.dtype)

                            loaded_params.add(name)
                        except AttributeError:
                            missing_params.append(name)
                            # Only keep existing non-meta parameters
                            if not param.is_meta:
                                new_state_dict[name] = param.data
                    else:
                        missing_params.append(name)
                        # Only keep existing non-meta parameters
                        if not param.is_meta:
                            new_state_dict[name] = param.data

    # Load the new state dict into the target model
    # Use strict=False to handle any mismatches gracefully
    target_model.load_state_dict(new_state_dict, assign=True)

    # Disable gradients for all parameters (inference mode)
    for param in target_model.parameters():
        param.requires_grad_(False)

    # Report loading status
    print(
        f"Successfully loaded {len(loaded_params)}/{len(list(target_model.named_parameters()))} parameters")
    if missing_params:
        print(
            f"Warning: Could not find {len(missing_params)} parameters in HF model:")
        for param in missing_params[:10]:  # Show first 10
            print(f"  - {param}")
        if len(missing_params) > 10:
            print(f"  ... and {len(missing_params) - 10} more")

    # Check if there are still meta parameters after loading
    meta_params = [name for name,
                   param in target_model.named_parameters() if param.is_meta]
    if meta_params:
        print(
            f"ERROR: {len(meta_params)} parameters are still on meta device after loading!")
        for param in meta_params[:5]:
            print(f"  - {param}")
        if len(meta_params) > 5:
            print(f"  ... and {len(meta_params) - 5} more")
        raise RuntimeError(
            f"Failed to materialize {len(meta_params)} meta parameters")