Implementation:Mlc ai Mlc llm Preshard

Overview

Preshard is a support module in MLC LLM that implements pre-sharding of model weight tensors for tensor parallelism. It is located at python/mlc_llm/support/preshard.py (124 lines).

Pre-sharding splits model weight parameters into multiple shards before deployment, so that each worker in a tensor-parallel inference setup can directly load its own shard without performing the split at runtime. The module uses TVM's Relax IR to generate and compile efficient sharding functions.

Purpose

In tensor-parallel LLM inference, large weight matrices must be distributed across multiple GPUs. Rather than loading full weights and splitting them at runtime, pre-sharding performs this split during weight conversion. This reduces memory requirements during loading and eliminates the runtime cost of weight distribution.

Architecture

The module operates in three phases:

1. IR Generation -- Builds TVM Relax IR functions that describe how to shard each parameter.
2. Compilation -- Compiles the generated IR into executable functions using DLight scheduling.
3. Application -- Returns callable shard functions alongside updated parameter names.

Core Functions

apply_preshard

def apply_preshard(
    named_params: Dict[str, nn.Parameter],
    tensor_parallel_shards: int,
    args: Any,
) -> Tuple[Dict[str, nn.Parameter], Dict[str, Callable[[Tensor], Sequence[Tensor]]]]:

The main entry point that processes all model parameters and produces sharding functions. The function:

1. Creates a relax.BlockBuilder for building the IR module.
2. Iterates over all named parameters, checking each for a shard_strategy attribute.
3. For parameters with a shard strategy:
   • Creates new parameter entries named {original_name}_shard-{worker_id} for each worker.
   • Records the mapping from original parameter name to its shard function name.
   • Calls _create_shard_func to generate the IR (avoiding duplicate generation for shared strategies).
4. For parameters without a shard strategy, passes them through unchanged.
5. If no parameters have shard strategies, logs a warning and continues without sharding.
6. Finalizes the IR module and compiles it via _compile_shard_funcs.
7. Replaces function name references with actual callable VM functions.

Returns: A tuple of:

• new_named_params -- Updated parameter dictionary with sharded names
• param_to_shard_func -- Mapping from original parameter names to callable shard functions

_sharded_param_name

def _sharded_param_name(param_name, worker_id):
    return f"{param_name}_shard-{worker_id}"

Generates the naming convention for sharded parameters (e.g., model.layers.0.self_attn.q_proj.weight_shard-0).

_create_shard_func

def _create_shard_func(
    bb: relax.BlockBuilder, param: nn.Parameter, tensor_parallel_shards: int
):

Generates a Relax IR function that implements the sharding logic for a single parameter. The generated function performs three operations:

1. Shard -- Calls the TIR (Tensor IR) shard function from the parameter's shard strategy, producing a tensor of shape [num_shards, *sharded_weight_shape].
2. Split -- Splits the result along axis 0 into num_shards individual tensors, each of shape [1, *sharded_weight_shape].
3. Squeeze -- Removes the leading dimension from each shard, yielding num_shards tensors of shape [*sharded_weight_shape].

The input weight shape is computed by scaling the shard dimension back to the full (unsharded) size:

weight_shape = param.shape
weight_shape[shard_strategy.dim] = weight_shape[shard_strategy.dim] * tensor_parallel_shards

_compile_shard_funcs

def _compile_shard_funcs(mod: IRModule, device: Device):

Compiles the IR module into executable functions using TVM's compilation pipeline:

1. Determines the target from the device.
2. Applies relax.transform.LegalizeOps to lower Relax operators to TIR.
3. Applies DLight auto-scheduling passes for GPU execution:
   • Matmul, GEMV, Reduction, GeneralReduction, Fallback
4. Builds the module with relax.build.
5. Creates a relax.VirtualMachine for execution.

with target:
    mod = relax.transform.LegalizeOps()(mod)
    mod = dl.ApplyDefaultSchedule(
        dl.gpu.Matmul(),
        dl.gpu.GEMV(),
        dl.gpu.Reduction(),
        dl.gpu.GeneralReduction(),
        dl.gpu.Fallback(),
    )(mod)
ex = relax.build(mod, target=target)
vm = relax.VirtualMachine(ex, device)

Warning Behavior

If no parameters in the model have a shard_strategy attribute, the module logs a warning:

logger.warning(
    "No parameters with 'shard_strategy' found."
    "At least one parameter must have a 'shard_strategy' for presharding. "
    "The model will continue to convert weights in a non-presharded manner."
)

The module still proceeds with compilation in this case, returning the original parameters unchanged.

Dependencies

• logging -- Standard library logging
• tvm.IRModule, tvm.relax -- TVM Relax IR for function building and compilation
• tvm.relax.frontend.nn -- Neural network parameter abstraction
• tvm.runtime.Device, tvm.runtime.Tensor -- TVM runtime types
• tvm.s_tir.dlight -- DLight auto-scheduling for GPU kernels
• tvm.target.Target -- Device target specification

File Location

python/mlc_llm/support/preshard.py