Aluode/PerceptionLabPortable

1# Copyright 2025 The HuggingFace Team. All rights reserved.2#3# Licensed under the Apache License, Version 2.0 (the "License");4# you may not use this file except in compliance with the License.5# You may obtain a copy of the License at6#7#     http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# distributed under the License is distributed on an "AS IS" BASIS,11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# See the License for the specific language governing permissions and13# limitations under the License.14 15from ..utils import is_accelerate_available, is_torch_available, logging16 17 18if is_torch_available():19    import torch20    from torch import nn21 22if is_accelerate_available():23    from accelerate import init_empty_weights24 25import re26from contextlib import contextmanager27 28 29logger = logging.get_logger(__name__)30 31FP4_VALUES = [32    +0.0,33    +0.5,34    +1.0,35    +1.5,36    +2.0,37    +3.0,38    +4.0,39    +6.0,40    -0.0,41    -0.5,42    -1.0,43    -1.5,44    -2.0,45    -3.0,46    -4.0,47    -6.0,48]49 50 51@contextmanager52def on_device(dev):53    if is_torch_available():54        import torch55 56        if isinstance(dev, torch.Tensor):57            dev = dev.device58        elif isinstance(dev, str):59            dev = torch.device(dev)60        dev_type = getattr(dev, "type", None)61        if dev_type == "cuda":62            with torch.cuda.device(dev):63                yield64                return65        if dev_type == "xpu" and hasattr(torch, "xpu"):66            with torch.xpu.device(dev):67                yield68                return69    # other: CPU70    yield71 72 73# Copied from GPT_OSS repo and vllm74def quantize_to_mxfp4(w, triton_kernels_hub):75    downcast_to_mxfp_torch = triton_kernels_hub.numerics_details.mxfp.downcast_to_mxfp_torch76    w, w_scale = downcast_to_mxfp_torch(w.to(torch.bfloat16), torch.uint8, axis=1)77    return w, w_scale78 79 80def swizzle_mxfp4(w, w_scale, triton_kernels_hub):81    """82    Changes the layout of the tensors depending on the hardware83    """84    FP4, convert_layout, wrap_torch_tensor = (85        triton_kernels_hub.tensor.FP4,86        triton_kernels_hub.tensor.convert_layout,87        triton_kernels_hub.tensor.wrap_torch_tensor,88    )89    layout = triton_kernels_hub.tensor_details.layout90    StridedLayout = triton_kernels_hub.tensor_details.layout.StridedLayout91 92    value_layout, value_layout_opts = layout.make_default_matmul_mxfp4_w_layout(mx_axis=1)93    w = convert_layout(wrap_torch_tensor(w, dtype=FP4), value_layout, **value_layout_opts)94    w_scale = convert_layout(wrap_torch_tensor(w_scale), StridedLayout)95    return w, w_scale96 97 98# Copied from GPT_OSS repo99# TODO: Add absolute link when the repo is public100def convert_moe_packed_tensors(101    blocks,102    scales,103    *,104    dtype: torch.dtype = torch.bfloat16,105    rows_per_chunk: int = 32768 * 1024,  # TODO these values are not here by mistake ;)106) -> torch.Tensor:107    """108    Convert the mxfp4 weights again, dequantizing and makes them compatible with the forward109    pass of GPT_OSS.110    """111    import math112 113    # Check if blocks and scales are on CPU, and move to GPU if so114    if not blocks.is_cuda and torch.cuda.is_available():115        blocks = blocks.cuda()116        scales = scales.cuda()117 118    scales = scales.to(torch.int32) - 127  # TODO that's because 128=2**7119 120    assert blocks.shape[:-1] == scales.shape, f"{blocks.shape[:-1]=} does not match {scales.shape=}"121 122    lut = torch.tensor(FP4_VALUES, dtype=dtype, device=blocks.device)123 124    *prefix_shape, G, B = blocks.shape125    rows_total = math.prod(prefix_shape) * G126 127    blocks = blocks.reshape(rows_total, B)128    scales = scales.reshape(rows_total, 1)129 130    out = torch.empty(rows_total, B * 2, dtype=dtype, device=blocks.device)131 132    for r0 in range(0, rows_total, rows_per_chunk):133        r1 = min(r0 + rows_per_chunk, rows_total)134 135        blk = blocks[r0:r1]136        exp = scales[r0:r1]137 138        # nibble indices -> int64139        idx_lo = (blk & 0x0F).to(torch.long)140        idx_hi = (blk >> 4).to(torch.long)141 142        sub = out[r0:r1]143        sub[:, 0::2] = lut[idx_lo]144        sub[:, 1::2] = lut[idx_hi]145 146        torch.ldexp(sub, exp, out=sub)147        del idx_lo, idx_hi, blk, exp, sub148 149    out = out.reshape(*prefix_shape, G, B * 2).view(*prefix_shape, G * B * 2)150    del blocks, scales, lut151    return out.transpose(1, 2).contiguous()152 153 154class Mxfp4GptOssExperts(nn.Module):155    def __init__(self, config):156        super().__init__()157 158        self.num_experts = config.num_local_experts159        self.intermediate_size = config.intermediate_size160        self.hidden_size = config.hidden_size161 162        self.gate_up_proj_blocks = nn.Parameter(163            torch.zeros(self.num_experts, 2 * self.intermediate_size, self.hidden_size // 32, 16, dtype=torch.uint8),164            requires_grad=False,165        )166        self.gate_up_proj_scales = nn.Parameter(167            torch.zeros(self.num_experts, 2 * self.intermediate_size, self.hidden_size // 32, dtype=torch.uint8),168            requires_grad=False,169        )170        self.gate_up_proj_bias = nn.Parameter(171            torch.zeros(self.num_experts, 2 * self.intermediate_size, dtype=torch.float32), requires_grad=False172        )173 174        self.down_proj_blocks = nn.Parameter(175            torch.zeros((self.num_experts, self.hidden_size, self.intermediate_size // 32, 16), dtype=torch.uint8),176            requires_grad=False,177        )178        self.down_proj_scales = nn.Parameter(179            torch.zeros(self.num_experts, self.hidden_size, self.intermediate_size // 32, dtype=torch.uint8),180            requires_grad=False,181        )182        self.down_proj_bias = nn.Parameter(183            torch.zeros(self.num_experts, self.hidden_size, dtype=torch.float32), requires_grad=False184        )185        self.alpha = 1.702186        self.limit = getattr(config, "swiglu_limit", 7.0)187        self.gate_up_proj_precision_config = None188        self.down_proj_precision_config = None189        self.limit = getattr(config, "swiglu_limit", 7.0)190 191    def forward(self, hidden_states: torch.Tensor, routing_data, gather_idx, scatter_idx) -> torch.Tensor:192        FnSpecs, FusedActivation, matmul_ogs = (193            triton_kernels_hub.matmul_ogs.FnSpecs,194            triton_kernels_hub.matmul_ogs.FusedActivation,195            triton_kernels_hub.matmul_ogs.matmul_ogs,196        )197        swiglu_fn = triton_kernels_hub.swiglu.swiglu_fn198 199        with on_device(hidden_states.device):200            act = FusedActivation(FnSpecs("swiglu", swiglu_fn, ("alpha", "limit")), (self.alpha, self.limit), 2)201 202            intermediate_cache1 = matmul_ogs(203                hidden_states,204                self.gate_up_proj,205                self.gate_up_proj_bias.to(torch.float32),206                routing_data,207                gather_indx=gather_idx,208                precision_config=self.gate_up_proj_precision_config,209                gammas=None,210                fused_activation=act,211            )212 213            intermediate_cache3 = matmul_ogs(214                intermediate_cache1,215                self.down_proj,216                self.down_proj_bias.to(torch.float32),217                routing_data,218                scatter_indx=scatter_idx,219                precision_config=self.down_proj_precision_config,220                gammas=routing_data.gate_scal,221            )222        return intermediate_cache3223 224 225# Adapted from GPT_OSS repo226# TODO: Add absolute link when the repo is public227def routing_torch_dist(228    logits,229    n_expts_act,230):231    import os232 233    GatherIndx, RoutingData, ScatterIndx, compute_expt_data_torch = (234        triton_kernels_hub.routing.GatherIndx,235        triton_kernels_hub.routing.RoutingData,236        triton_kernels_hub.routing.ScatterIndx,237        triton_kernels_hub.routing.compute_expt_data_torch,238    )239 240    with on_device(logits.device):241        world_size = torch.distributed.get_world_size()242        rank = int(os.environ.get("LOCAL_RANK", "0"))243        replace_value = -1244 245        n_tokens = logits.shape[0]246        n_expts_tot = logits.shape[1]247 248        n_local_experts = n_expts_tot // world_size249        local_expert_start = rank * n_local_experts250        local_expert_end = (rank + 1) * n_local_experts251 252        n_gates_pad = n_tokens * n_expts_act253 254        def topk(vals, k):255            tk_indx = torch.argsort(-vals, dim=1, stable=True)[:, :k]256            tk_indx = tk_indx.long()257            tk_val = torch.take_along_dim(vals, tk_indx, dim=1)258            return tk_val, tk_indx.int()259 260        expt_scal, expt_indx = topk(logits, n_expts_act)261        expt_scal = torch.softmax(expt_scal, dim=-1)262        expt_indx, sort_indices = torch.sort(expt_indx, dim=1)263        expt_scal = torch.gather(expt_scal, 1, sort_indices)264 265        # Flatten and mask for local experts266        expt_scal = expt_scal.reshape(-1)267 268        hist = torch.histc(expt_indx, bins=n_expts_tot, max=n_expts_tot - 1)[local_expert_start:local_expert_end]269 270        expt_indx = expt_indx.view(-1).to(torch.int32)271 272        # we use a large value to replace the indices that are not in the local expert range273        var = 1000274        expt_indx = torch.where(expt_indx < local_expert_start, var, expt_indx)275        topk_indx = torch.argsort(expt_indx, stable=True).to(torch.int32)276        gate_indx = torch.argsort(topk_indx).to(torch.int32)277        expt_indx = torch.where(expt_indx < local_expert_end, expt_indx, replace_value)278        expt_indx = torch.where(local_expert_start <= expt_indx, expt_indx, replace_value)279 280        gate_indx = torch.where(expt_indx == replace_value, replace_value, gate_indx)281        gate_scal = expt_scal[topk_indx]282 283        topk_indx = torch.where(gate_indx[topk_indx] == replace_value, replace_value, topk_indx)284 285        # # Routing metadata for local expert computation286        gather_indx = GatherIndx(src_indx=topk_indx.int(), dst_indx=gate_indx.int())287        scatter_indx = ScatterIndx(src_indx=gate_indx.int(), dst_indx=topk_indx.int())288 289        expt_data = compute_expt_data_torch(hist, n_local_experts, n_gates_pad)290 291        hit_experts = n_expts_act292    return RoutingData(gate_scal, hist, n_local_experts, hit_experts, expt_data), gather_indx, scatter_indx293 294 295def mlp_forward(self, hidden_states):296    import torch.distributed as dist297 298    if dist.is_available() and dist.is_initialized() and hasattr(self, "_is_hooked"):299        routing = routing_torch_dist300    else:301        routing = triton_kernels_hub.routing.routing302 303    batch_size = hidden_states.shape[0]304    hidden_states = hidden_states.reshape(-1, self.router.hidden_dim)305    router_logits = nn.functional.linear(hidden_states, self.router.weight, self.router.bias)306 307    with on_device(router_logits.device):308        routing_data, gather_idx, scatter_idx = routing(router_logits, self.router.top_k)309 310    routed_out = self.experts(hidden_states, routing_data, gather_idx, scatter_idx)311    routed_out = routed_out.reshape(batch_size, -1, self.router.hidden_dim)312    return routed_out, router_logits313 314 315def should_convert_module(current_key_name, patterns):316    current_key_name_str = ".".join(current_key_name)317    if not any(318        re.match(f"{key}\\.", current_key_name_str) or re.match(f"{key}", current_key_name_str) for key in patterns319    ):320        return True321    return False322 323 324def dequantize(module, param_name, param_value, target_device, dq_param_name, **kwargs):325    from ..integrations.tensor_parallel import shard_and_distribute_module326 327    model = kwargs.get("model")328    empty_param = kwargs.get("empty_param")329    casting_dtype = kwargs.get("casting_dtype")330    to_contiguous = kwargs.get("to_contiguous")331    rank = kwargs.get("rank")332    device_mesh = kwargs.get("device_mesh")333 334    for proj in ["gate_up_proj", "down_proj"]:335        if proj in param_name:336            if device_mesh is not None:337                param_value = shard_and_distribute_module(338                    model,339                    param_value,340                    empty_param,341                    dq_param_name,342                    casting_dtype,343                    to_contiguous,344                    rank,345                    device_mesh,346                )347            blocks_attr = f"{proj}_blocks"348            scales_attr = f"{proj}_scales"349            setattr(module, param_name.rsplit(".", 1)[1], param_value)350            if hasattr(module, blocks_attr) and hasattr(module, scales_attr):351                dequantized = convert_moe_packed_tensors(getattr(module, blocks_attr), getattr(module, scales_attr))352                if target_device == "cpu" and torch.cuda.is_available():353                    torch.cuda.empty_cache()354                setattr(module, proj, torch.nn.Parameter(dequantized.to(target_device)))355                delattr(module, blocks_attr)356                delattr(module, scales_attr)357 358 359def load_and_swizzle_mxfp4(module, param_name, param_value, target_device, triton_kernels_hub, **kwargs):360    """361    This transforms the weights obtained using `convert_gpt_oss.py` to load them into `Mxfp4GptOssExperts`.362    """363    PrecisionConfig, FlexCtx, InFlexData = (364        triton_kernels_hub.matmul_ogs.PrecisionConfig,365        triton_kernels_hub.matmul_ogs.FlexCtx,366        triton_kernels_hub.matmul_ogs.InFlexData,367    )368    from ..integrations.tensor_parallel import shard_and_distribute_module369 370    model = kwargs.get("model")371    empty_param = kwargs.get("empty_param")372    casting_dtype = kwargs.get("casting_dtype")373    to_contiguous = kwargs.get("to_contiguous")374    rank = kwargs.get("rank")375    device_mesh = kwargs.get("device_mesh")376    if "blocks" in param_name:377        proj = param_name.split(".")[-1].split("_blocks")[0]378    if "scales" in param_name:379        proj = param_name.split(".")[-1].split("_scales")[0]380    if device_mesh is not None:381        shard_and_distribute_module(382            model, param_value, empty_param, param_name, casting_dtype, to_contiguous, rank, device_mesh383        )384    else:385        setattr(module, param_name.rsplit(".", 1)[1], torch.nn.Parameter(param_value, requires_grad=False))386    blocks_attr = f"{proj}_blocks"387    scales_attr = f"{proj}_scales"388    blocks = getattr(module, blocks_attr)  # at this point values were loaded from ckpt389    scales = getattr(module, scales_attr)390    # Check if both blocks and scales both not on meta device391    if blocks.device.type != "meta" and scales.device.type != "meta":392        local_experts = blocks.size(0)393        if proj == "gate_up_proj":394            blocks = blocks.reshape(local_experts, module.intermediate_size * 2, -1)395        else:396            blocks = blocks.reshape(local_experts, -1, module.intermediate_size // 2)397        if getattr(target_device, "type", target_device) == "cpu":398            target_device = "cuda"399        blocks = blocks.to(target_device).contiguous()400        scales = scales.to(target_device).contiguous()401        with on_device(target_device):402            triton_weight_tensor, weight_scale = swizzle_mxfp4(403                blocks.transpose(-2, -1), scales.transpose(-2, -1), triton_kernels_hub404            )405 406        # need to overwrite the shapes for the kernels407        if proj == "gate_up_proj":408            triton_weight_tensor.shape = torch.Size([local_experts, module.hidden_size, module.intermediate_size * 2])409        else:410            triton_weight_tensor.shape = torch.Size([local_experts, module.intermediate_size, module.hidden_size])411 412        # triton_weight_tensor is what needs to be passed in oai kernels. It stores the data, the shapes and any more objects. It is like a subtensor413        setattr(module, proj, triton_weight_tensor)414        setattr(415            module,416            f"{proj}_precision_config",417            PrecisionConfig(weight_scale=weight_scale, flex_ctx=FlexCtx(rhs_data=InFlexData())),418        )419 420        # delete blocks and scales421        delattr(module, scales_attr)422        delattr(module, blocks_attr)423        del blocks424 425 426def _replace_with_mxfp4_linear(427    model,428    modules_to_not_convert=None,429    current_key_name=None,430    quantization_config=None,431    has_been_replaced=False,432    config=None,433):434    if current_key_name is None:435        current_key_name = []436 437    for name, module in model.named_children():438        current_key_name.append(name)439        if not should_convert_module(current_key_name, modules_to_not_convert):440            current_key_name.pop(-1)441            continue442        if module.__class__.__name__ == "GptOssExperts" and not quantization_config.dequantize:443            with init_empty_weights():444                model._modules[name] = Mxfp4GptOssExperts(config)445                has_been_replaced = True446        if module.__class__.__name__ == "GptOssMLP" and not quantization_config.dequantize:447            from types import MethodType448 449            module.forward = MethodType(mlp_forward, module)450        if len(list(module.children())) > 0:451            _, has_been_replaced = _replace_with_mxfp4_linear(452                module,453                modules_to_not_convert,454                current_key_name,455                quantization_config,456                has_been_replaced=has_been_replaced,457                config=config,458            )459        current_key_name.pop(-1)460    return model, has_been_replaced461 462 463def replace_with_mxfp4_linear(464    model,465    modules_to_not_convert=None,466    current_key_name=None,467    quantization_config=None,468    config=None,469):470    if quantization_config.dequantize:471        return model472    else:473        from kernels import get_kernel474 475        global triton_kernels_hub476        triton_kernels_hub = get_kernel("kernels-community/triton_kernels")477 478    modules_to_not_convert = ["lm_head"] if modules_to_not_convert is None else modules_to_not_convert479 480    if quantization_config.modules_to_not_convert is not None:481        modules_to_not_convert.extend(quantization_config.modules_to_not_convert)482    modules_to_not_convert = list(set(modules_to_not_convert))483    model, has_been_replaced = _replace_with_mxfp4_linear(484        model,485        modules_to_not_convert,486        current_key_name,487        quantization_config,488        config=config,489    )490    if not has_been_replaced:491        logger.warning(492            "You are loading your model using mixed-precision FP4 quantization but no linear modules were found in your model."493            " Please double check your model architecture, or submit an issue on github if you think this is"494            " a bug."495        )496 497    return model498