def shard_low_precision_checkpoint(
    low_precision_checkpoint,
    model_config,
    rank,
    world_size,
    quantization_method,
    tp_grain_size,
    desc_act,
    bits,
):
    r"""
    Shard low-precision checkpoint for autoTP.
    Sharding policy:
        Attention:
            Divide N or K by num_k_v_heads if available or num_attention_heads
            Assign each rank a chunk of N or K evenly or by the 2:1:1 policy
            - 2:1:1 policy: 4 heads for 3 ranks, the first rank gets 2 and others get 1.
        MLP:
            Divide N or K by tp_grain_size
            Assign each rank a chunk of N or K evenly or by the 2:1:1 policy
    Tensor shapes for different quantization methods:
        AWQ:
            qweight shape: [K, N // 8]
            scales shape: [K // G, N]
            qzeros shape: [K // G, N // 8]
        GPTQ:
            qweight shape: [K // 8, N]
            scales shape: [K // G, N]
            qzeros shape: [K // G, N // 8]
            g_idx shape: [K]
            Similar to AWQ, but with different paths if desc_act is True or False.
            If desc_act is True:
            - If sharding along N
                - g_idx is not sharded
                - Others are sharded as usual
            - If sharding along K
                - g_idx is shared the same way as qweight
                - Scales and zero points are not sharded
            If desc_act is False:
            - g_idx is ignored
            - Others are sharded as usual
        FP8:
            qweight shape: [N, K]
            weight_scale_inv shape: [N // block_n, K // block_k], block_n = block_k = 128 for DeepSeek-V3/R1
        RTN (INC):
            qweight shape: [K // 4, N]
            scales shape: [K // G, N] or [1, N]
            qzeros shape: [K // G, N // 4] or [1, N]
        INT8 (meituan/DeepSeek-R1-Channel-INT8):
            qweight shape: [N, K]
            scales shape: [N, 1]
            qzeros shape: None

    Args:
        low_precision_checkpoint (dict): Model's low_precision_checkpoint as a state dict.
        model_config (dict): HuggingFace model config as a dict.
        rank (int): current rank for tp.
        world_size (int): total number of ranks.
        quantization_method (str): "awq" or "gptq".
        tp_grain_size (int): Grain size for MLP sharding. Usually equal to group size for
            quantization. Must be a multiple of 8.
        desc_act (bool): desc_act (a.k.a. act_order) for GPTQ. False for others.
        bits: Number of bits for quantization.

    Returns:
        A sharded checkpoint dict.

    """
    assert tp_grain_size % 8 == 0, "tp_grain_size must be a multiple of 8"
    if isinstance(model_config, dict):
        num_heads = model_config["num_attention_heads"]
        if "num_key_value_heads" in model_config:
            num_heads = model_config["num_key_value_heads"]
    else:
        num_heads = model_config.num_attention_heads
        if "num_key_value_heads" in model_config:
            num_heads = model_config.num_key_value_heads
    local_rank = rank

    mha_layers_split_by_N = [
        "q_proj",
        "k_proj",
        "v_proj",
        "q_b_proj",
        "kv_b_proj",
    ]
    # mlp is split with grain size = tp_grain_size
    mlp_layers_split_by_N = [
        "gate_proj",
        "up_proj",
        "fc_in",
        "fc1",
        "query_key_value",
        "w1",
        "w3",
    ]
    mha_layers_split_by_K = [
        "o_proj",
        "out_proj",
    ]
    # mlp is split with grain size = tp_grain_size
    mlp_layers_split_by_K = [
        "down_proj",
        "fc_out",
        "fc2",
        "dense",
        "dense_4h_to_h",
        "w2",
    ]
    lm_head_layers = ["lm_head"]  # split by K but not quantized
    low_precision_checkpoint_dict = low_precision_checkpoint.copy()
    head_range = [0]
    head_per_rank = num_heads // world_size
    for i in range(0, world_size):
        head_this_rank = head_per_rank
        if i < num_heads % world_size:
            head_this_rank += 1
        head_range.append(head_range[-1] + head_this_rank)
    for key in low_precision_checkpoint.keys():
        q_head_start = head_range[rank]
        q_head_end = q_head_start + (head_range[rank + 1] - head_range[rank])
        if "bias" in key:
            continue
        if any(substring in key for substring in mha_layers_split_by_N):
            data = low_precision_checkpoint_dict[key]
            if quantization_method == "awq":
                # qweight shape: [K, N // 8]
                # scales shape: [K // G, N]
                # qzeros shape: [K // G, N // 8]
                if data.shape[-1] % head_range[-1] == 0:
                    dim = data.shape[-1] // head_range[-1]
                else:
                    assert data.shape[-1] % world_size == 0
                    dim = data.shape[-1] // world_size
                    q_head_start = local_rank
                    q_head_end = local_rank + 1
                low_precision_checkpoint_dict[key] = data[
                    :, q_head_start * dim : q_head_end * dim
                ].contiguous()
            elif quantization_method == "gptq" or (
                quantization_method == "rtn" and bits == 4
            ):
                # qweight shape: [K // 8, N]
                # scales shape: [K // G, N]
                # qzeros shape: [K // G, N // 8]
                # g_idx shape: [K]
                if data.shape[-1] % head_range[-1] == 0:
                    dim = data.shape[-1] // head_range[-1]
                else:
                    assert data.shape[-1] % world_size == 0
                    dim = data.shape[-1] // world_size
                    q_head_start = local_rank
                    q_head_end = local_rank + 1
                if "g_idx" in key:
                    if not desc_act:
                        low_precision_checkpoint_dict.pop(key)
                else:
                    low_precision_checkpoint_dict[key] = data[
                        :, q_head_start * dim : q_head_end * dim
                    ].contiguous()
            elif quantization_method == "fp8":
                # weight shape: [N, K]
                # weight_scale_inv shape: [N // block_n, K // block_k]
                if data.shape[0] % head_range[-1] == 0:
                    dim = data.shape[0] // head_range[-1]
                else:
                    assert data.shape[0] % world_size == 0
                    dim = data.shape[0] // world_size
                    q_head_start = local_rank
                    q_head_end = local_rank + 1
                low_precision_checkpoint_dict[key] = data[
                    q_head_start * dim : q_head_end * dim, :
                ].contiguous()
            elif quantization_method == "rtn" and bits == 8:
                # from INC, using GPTQ-like format
                # qweight shape: [K // 4, N]
                # scales shape: [K // G, N] or [1, N]
                # qzeros shape: [K // G, N // 4] or [1, N // 4]
                if data.shape[-1] % head_range[-1] == 0:
                    dim = data.shape[-1] // head_range[-1]
                else:
                    assert data.shape[-1] % world_size == 0
                    dim = data.shape[-1] // world_size
                    q_head_start = local_rank
                    q_head_end = local_rank + 1
                low_precision_checkpoint_dict[key] = data[
                    :, q_head_start * dim : q_head_end * dim
                ].contiguous()
            elif quantization_method == "int8":
                # qweight shape: [N, K]
                # scales shape: [N, 1]
                # qzeros shape: None
                if data.shape[0] % head_range[-1] == 0:
                    dim = data.shape[0] // head_range[-1]
                else:
                    assert data.shape[0] % world_size == 0
                    dim = data.shape[0] // world_size
                    q_head_start = local_rank
                    q_head_end = local_rank + 1
                low_precision_checkpoint_dict[key] = data[
                    q_head_start * dim : q_head_end * dim, :
                ].contiguous()
            else:
                raise AssertionError(f"{quantization_method} is not supported yet.")
        elif any(substring in key for substring in mlp_layers_split_by_N):
            data = low_precision_checkpoint_dict[key]
            if quantization_method == "awq":
                # qweight shape: [K, N // 8]
                # scales shape: [K // G, N]
                # qzeros shape: [K // G, N // 8]
                if "scales" in key:
                    assert (
                        data.shape[1] % tp_grain_size == 0
                    ), "N must be divisible by tp_grain_size"
                    grains = data.shape[1] // tp_grain_size
                    dim = tp_grain_size
                else:
                    assert (
                        data.shape[1] * 8
                    ) % tp_grain_size == 0, "N must be divisible by tp_grain_size"
                    grains = data.shape[1] // (tp_grain_size // 8)
                    dim = tp_grain_size // 8
                grains_per_rank = grains // world_size
                grains_rem = grains % world_size
                grains_start = grains_per_rank * local_rank + min(
                    local_rank, grains_rem
                )
                grains_end = (
                    grains_start
                    + grains_per_rank
                    + (1 if local_rank < grains_rem else 0)
                )
                low_precision_checkpoint_dict[key] = data[
                    :, grains_start * dim : grains_end * dim
                ].contiguous()
            elif quantization_method == "gptq" or (
                quantization_method == "rtn" and bits == 4
            ):
                # qweight shape: [K // 8, N]
                # scales shape: [K // G, N]
                # qzeros shape: [K // G, N // 8]
                # g_idx shape: [K]
                if "qzeros" in key:
                    assert (
                        data.shape[-1] * 8
                    ) % tp_grain_size == 0, "N must be divisible by tp_grain_size"
                    grains = data.shape[-1] // (tp_grain_size // 8)
                    dim = tp_grain_size // 8
                elif "g_idx" not in key:  # qweight, scales
                    assert (
                        data.shape[-1] % tp_grain_size == 0
                    ), "N must be divisible by tp_grain_size"
                    grains = data.shape[-1] // tp_grain_size
                    dim = tp_grain_size
                grains_per_rank = grains // world_size
                grains_rem = grains % world_size
                grains_start = grains_per_rank * local_rank + min(
                    local_rank, grains_rem
                )
                grains_end = (
                    grains_start
                    + grains_per_rank
                    + (1 if local_rank < grains_rem else 0)
                )
                if "g_idx" in key:
                    if not desc_act:
                        low_precision_checkpoint_dict.pop(key)
                else:
                    low_precision_checkpoint_dict[key] = data[
                        :, grains_start * dim : grains_end * dim
                    ].contiguous()
            elif quantization_method == "fp8":
                # weight shape: [N, K]
                # weight_scale_inv shape: [N // block_n, K // block_k]
                assert (
                    data.shape[0] % tp_grain_size == 0
                ), "N must be divisible by tp_grain_size"
                grains = data.shape[0] // tp_grain_size
                dim = tp_grain_size
                grains_per_rank = grains // world_size
                grains_rem = grains % world_size
                grains_start = grains_per_rank * local_rank + min(
                    local_rank, grains_rem
                )
                grains_end = (
                    grains_start
                    + grains_per_rank
                    + (1 if local_rank < grains_rem else 0)
                )
                low_precision_checkpoint_dict[key] = data[
                    grains_start * dim : grains_end * dim, :
                ].contiguous()
            elif quantization_method == "rtn" and bits == 8:
                # from INC, using GPTQ-like format
                # qweight shape: [K // 4, N]
                # scales shape: [K // G, N] or [1, N]
                # qzeros shape: [K // G, N // 4] or [1, N // 4]
                comp_ratio = 4
                if "qzeros" in key:
                    assert (
                        data.shape[-1] * comp_ratio
                    ) % tp_grain_size == 0, "N must be divisible by tp_grain_size"
                    grains = data.shape[-1] // (tp_grain_size // comp_ratio)
                    dim = tp_grain_size // comp_ratio
                grains_per_rank = grains // world_size
                grains_rem = grains % world_size
                grains_start = grains_per_rank * local_rank + min(
                    local_rank, grains_rem
                )
                grains_end = (
                    grains_start
                    + grains_per_rank
                    + (1 if local_rank < grains_rem else 0)
                )
                low_precision_checkpoint_dict[key] = data[
                    :, grains_start * dim : grains_end * dim
                ].contiguous()
            elif quantization_method == "int8":
                # qweight shape: [N, K]
                # scales shape: [N, 1]
                # qzeros shape: None
                assert (
                    data.shape[0] % tp_grain_size == 0
                ), "N must be divisible by tp_grain_size"
                grains = data.shape[0] // tp_grain_size
                dim = tp_grain_size
                grains_per_rank = grains // world_size
                grains_rem = grains % world_size
                grains_start = grains_per_rank * local_rank + min(
                    local_rank, grains_rem
                )
                grains_end = (
                    grains_start
                    + grains_per_rank
                    + (1 if local_rank < grains_rem else 0)
                )
                low_precision_checkpoint_dict[key] = data[
                    grains_start * dim : grains_end * dim, :
                ].contiguous()
            else:
                raise AssertionError(f"{quantization_method} is not supported yet.")
        elif any(substring in key for substring in mha_layers_split_by_K):
            data = low_precision_checkpoint_dict[key]
            if ("scales" in key or "qzeros" in key) and data.shape[0] == 1:
                continue
            if quantization_method == "awq":
                # qweight shape: [K, N // 8]
                # scales shape: [K // G, N]
                # qzeros shape: [K // G, N // 8]
                if data.shape[0] % head_range[-1] == 0:
                    dim = data.shape[0] // head_range[-1]
                else:
                    assert data.shape[0] % world_size == 0
                    dim = data.shape[0] // world_size
                    q_head_start = local_rank
                    q_head_end = local_rank + 1
                low_precision_checkpoint_dict[key] = data[
                    q_head_start * dim : q_head_end * dim
                ].contiguous()
            elif quantization_method == "gptq" or (
                quantization_method == "rtn" and bits == 4
            ):
                # qweight shape: [K // 8, N]
                # scales shape: [K // G, N]
                # qzeros shape: [K // G, N // 8]
                # g_idx shape: [K]
                if data.shape[0] % head_range[-1] == 0:
                    dim = data.shape[0] // head_range[-1]
                else:
                    assert data.shape[0] % world_size == 0
                    dim = data.shape[0] // world_size
                    q_head_start = local_rank
                    q_head_end = local_rank + 1
                if desc_act is False:
                    if "g_idx" in key:
                        low_precision_checkpoint_dict.pop(key)
                    else:
                        low_precision_checkpoint_dict[key] = data[
                            q_head_start * dim : q_head_end * dim
                        ].contiguous()
                elif "g_idx" in key or "qweight" in key:
                    low_precision_checkpoint_dict[key] = data[
                        q_head_start * dim : q_head_end * dim
                    ].contiguous()
            elif quantization_method == "fp8":
                # weight shape: [N, K]
                # weight_scale_inv shape: [N // block_n, K // block_k]
                if data.shape[-1] % head_range[-1] == 0:
                    dim = data.shape[-1] // head_range[-1]
                else:
                    assert data.shape[-1] % world_size == 0
                    dim = data.shape[-1] // world_size
                    q_head_start = local_rank
                    q_head_end = local_rank + 1
                low_precision_checkpoint_dict[key] = data[
                    :, q_head_start * dim : q_head_end * dim
                ].contiguous()
            elif quantization_method == "rtn" and bits == 8:
                # from INC, using GPTQ-like format
                # qweight shape: [K // 4, N]
                # scales shape: [K // G, N] or [1, N]
                # qzeros shape: [K // G, N // 4] or [1, N // 4]
                if data.shape[0] == 1:
                    continue
                if data.shape[0] % head_range[-1] == 0:
                    dim = data.shape[0] // head_range[-1]
                else:
                    assert data.shape[0] % world_size == 0
                    dim = data.shape[0] // world_size
                    q_head_start = local_rank
                    q_head_end = local_rank + 1
                low_precision_checkpoint_dict[key] = data[
                    q_head_start * dim : q_head_end * dim
                ].contiguous()
            elif quantization_method == "int8":
                # qweight shape: [N, K]
                # scales shape: [N, 1]
                # qzeros shape: None
                if data.shape[-1] == 1:
                    continue
                if data.shape[-1] % head_range[-1] == 0:
                    dim = data.shape[-1] // head_range[-1]
                else:
                    assert data.shape[-1] % world_size == 0
                    dim = data.shape[-1] // world_size
                    q_head_start = local_rank
                    q_head_end = local_rank + 1
                low_precision_checkpoint_dict[key] = data[
                    :, q_head_start * dim : q_head_end * dim
                ].contiguous()
            else:
                raise AssertionError(f"{quantization_method} is not supported yet.")
        elif any(substring in key for substring in mlp_layers_split_by_K):
            data = low_precision_checkpoint_dict[key]
            if ("scales" in key or "qzeros" in key) and data.shape[0] == 1:
                continue
            if quantization_method == "awq":
                # qweight shape: [K, N // 8]
                # scales shape: [K // G, N]
                # qzeros shape: [K // G, N // 8]
                if "qweight" in key:
                    assert (
                        data.shape[0] % tp_grain_size == 0
                    ), "K must be divisible by tp_grain_size"
                    grains = data.shape[0] // tp_grain_size
                    dim = tp_grain_size
                else:
                    grains = data.shape[0]
                    dim = 1
                grains_per_rank = grains // world_size
                grains_rem = grains % world_size
                grains_start = grains_per_rank * local_rank + min(
                    local_rank, grains_rem
                )
                grains_end = (
                    grains_start
                    + grains_per_rank
                    + (1 if local_rank < grains_rem else 0)
                )
                low_precision_checkpoint_dict[key] = data[
                    grains_start * dim : grains_end * dim
                ].contiguous()
            elif quantization_method == "gptq" or (
                quantization_method == "rtn" and bits == 4
            ):
                # qweight shape: [K // 8, N]
                # scales shape: [K // G, N]
                # qzeros shape: [K // G, N // 8]
                # g_idx shape: [K]
                if "qweight" in key:
                    assert (
                        data.shape[0] * 8 % tp_grain_size == 0
                    ), "K must be divisible by tp_grain_size"
                    grains = data.shape[0] // (tp_grain_size // 8)
                    dim = tp_grain_size // 8
                elif "g_idx" in key:
                    assert data.shape[0] % tp_grain_size == 0
                    grains = data.shape[0] // tp_grain_size
                    dim = tp_grain_size
                else:
                    grains = data.shape[0]
                    dim = 1
                grains_per_rank = grains // world_size
                grains_rem = grains % world_size
                grains_start = grains_per_rank * local_rank + min(
                    local_rank, grains_rem
                )
                grains_end = (
                    grains_start
                    + grains_per_rank
                    + (1 if local_rank < grains_rem else 0)
                )
                if desc_act is False:
                    if "g_idx" in key:
                        low_precision_checkpoint_dict.pop(key)
                    else:
                        low_precision_checkpoint_dict[key] = data[
                            grains_start * dim : grains_end * dim
                        ].contiguous()
                elif "g_idx" in key or "qweight" in key:
                    low_precision_checkpoint_dict[key] = data[
                        grains_start * dim : grains_end * dim
                    ].contiguous()
            elif quantization_method == "fp8":
                # weight shape: [N, K]
                # weight_scale_inv shape: [N // block_n, K // block_k]
                if "weight" in key and "scale" not in key:
                    assert (
                        data.shape[-1] % tp_grain_size == 0
                    ), "K must be divisible by tp_grain_size"
                    grains = data.shape[-1] // tp_grain_size
                    dim = tp_grain_size
                else:
                    grains = data.shape[-1]
                    dim = 1
                grains_per_rank = grains // world_size
                grains_rem = grains % world_size
                grains_start = grains_per_rank * local_rank + min(
                    local_rank, grains_rem
                )
                grains_end = (
                    grains_start
                    + grains_per_rank
                    + (1 if local_rank < grains_rem else 0)
                )
                low_precision_checkpoint_dict[key] = data[
                    :, grains_start * dim : grains_end * dim
                ].contiguous()
            elif quantization_method == "rtn" and bits == 8:
                # from INC, using GPTQ-like format
                # qweight shape: [K // 4, N]
                # scales shape: [K // G, N] or [1, N]
                # qzeros shape: [K // G, N // 4] or [1, N // 4]
                if data.shape[0] == 1:
                    continue
                comp_ratio = 4
                if "qweight" in key:
                    assert (
                        data.shape[0] * comp_ratio % tp_grain_size == 0
                    ), "K must be divisible by tp_grain_size"
                    grains = data.shape[0] // (tp_grain_size // comp_ratio)
                    dim = tp_grain_size // comp_ratio
                else:
                    grains = data.shape[0]
                    dim = 1
                grains_per_rank = grains // world_size
                grains_rem = grains % world_size
                grains_start = grains_per_rank * local_rank + min(
                    local_rank, grains_rem
                )
                grains_end = (
                    grains_start
                    + grains_per_rank
                    + (1 if local_rank < grains_rem else 0)
                )
                low_precision_checkpoint_dict[key] = data[
                    grains_start * dim : grains_end * dim
                ].contiguous()
            elif quantization_method == "int8":
                # qweight shape: [N, K]
                # scales shape: [N, 1]
                # qzeros shape: None
                if data.shape[-1] == 1:
                    continue
                if "scales" in key:
                    grains = data.shape[-1]
                    dim = 1
                else:
                    grains = data.shape[-1] // tp_grain_size
                    dim = tp_grain_size
                grains_per_rank = grains // world_size
                grains_rem = grains % world_size
                grains_start = grains_per_rank * local_rank + min(
                    local_rank, grains_rem
                )
                grains_end = (
                    grains_start
                    + grains_per_rank
                    + (1 if local_rank < grains_rem else 0)
                )
                low_precision_checkpoint_dict[key] = data[
                    :, grains_start * dim : grains_end * dim
                ].contiguous()
            else:
                raise AssertionError(f"{quantization_method} is not supported yet.")
        elif any(substring in key for substring in lm_head_layers):
            # lm_head: [N, K] (not quantized)
            # Same for all quantization methods
            data = low_precision_checkpoint_dict[key]
            assert (
                data.shape[1] % tp_grain_size == 0
            ), "K must be divisible by tp_grain_size"
            grains = data.shape[1] // tp_grain_size
            dim = tp_grain_size
            grains_per_rank = grains // world_size
            grains_rem = grains % world_size
            grains_start = grains_per_rank * local_rank + min(local_rank, grains_rem)
            grains_end = (
                grains_start + grains_per_rank + (1 if local_rank < grains_rem else 0)
            )
            low_precision_checkpoint_dict[key] = data[
                :, grains_start * dim : grains_end * dim
            ].contiguous()

    return low_precision_checkpoint_dict

    Shard low-precision checkpoint for autoTP.
    Sharding policy:
        Attention:
            Divide N or K by num_k_v_heads if available or num_attention_heads
            Assign each rank a chunk of N or K evenly or by the 2:1:1 policy
            - 2:1:1 policy: 4 heads for 3 ranks, the first rank gets 2 and others get 1.
        MLP:
            Divide N or K by tp_grain_size
            Assign each rank a chunk of N or K evenly or by the 2:1:1 policy
    Tensor shapes for different quantization methods:
        AWQ:
            qweight shape: [K, N // 8]
            scales shape: [K // G, N]
            qzeros shape: [K // G, N // 8]
        GPTQ:
            qweight shape: [K // 8, N]
            scales shape: [K // G, N]
            qzeros shape: [K // G, N // 8]
            g_idx shape: [K]
            Similar to AWQ, but with different paths if desc_act is True or False.
            If desc_act is True:
            - If sharding along N
                - g_idx is not sharded
                - Others are sharded as usual
            - If sharding along K
                - g_idx is shared the same way as qweight
                - Scales and zero points are not sharded
            If desc_act is False:
            - g_idx is ignored
            - Others are sharded as usual
        FP8:
            qweight shape: [N, K]
            weight_scale_inv shape: [N // block_n, K // block_k], block_n = block_k = 128 for DeepSeek-V3/R1
        RTN (INC):
            qweight shape: [K // 4, N]
            scales shape: [K // G, N] or [1, N]
            qzeros shape: [K // G, N // 4] or [1, N]
        INT8 (meituan/DeepSeek-R1-Channel-INT8):
            qweight shape: [N, K]
            scales shape: [N, 1]
            qzeros shape: None

    Args:
        low_precision_checkpoint (dict): Model's low_precision_checkpoint as a state dict.
        model_config (dict): HuggingFace model config as a dict.
        rank (int): current rank for tp.
        world_size (int): total number of ranks.
        quantization_method (str): "awq" or "gptq".
        tp_grain_size (int): Grain size for MLP sharding. Usually equal to group size for
            quantization. Must be a multiple of 8.
        desc_act (bool): desc_act (a.k.a. act_order) for GPTQ. False for others.
        bits: Number of bits for quantization.

    Returns:
        A sharded checkpoint dict.

    
intel_extension_for_pytorch/llm/utils.py