Implementation:Huggingface Optimum Initialize Parameter Meta

Overview

Attaches ParameterMeta metadata objects to each parameter in the model, detecting tied (shared) parameters and preparing the metadata structure for downstream parallelization passes.

Source

Property	Value
Function File	`optimum/fx/parallelization/utils.py`
Function Lines	L297-313
Dataclass File	`optimum/fx/parallelization/core.py`
Dataclass Lines	L69-103

APIs

initialize_parameter_meta

# optimum/fx/parallelization/utils.py L297-313
def initialize_parameter_meta(model: nn.Module) -> None:
    """Attach ParameterMeta to every parameter in the model.

    Detects tied parameters and marks them accordingly.
    """
    ...

ParameterMeta Dataclass

# optimum/fx/parallelization/core.py L69-103
@dataclass
class ParameterMeta:
    is_tied: bool = False
    is_parallel: bool = False
    is_modified_meta: bool = False
    need_initialize: bool = False
    init_fn: Optional[Callable] = None
    dim: int = 0
    mapping: Dict[HashableSlice, ParameterSlice] = field(default_factory=dict)

Import

from optimum.fx.parallelization.utils import initialize_parameter_meta
from optimum.fx.parallelization.core import ParameterMeta

ParameterMeta Fields

Field	Type	Default	Description
is_tied	`bool`	`False`	Whether this parameter is shared (tied) across multiple modules.
is_parallel	`bool`	`False`	Whether this parameter has been parallelized by a replacement pass.
is_modified_meta	`bool`	`False`	Whether the parameter's meta information has been modified from the original.
need_initialize	`bool`	`False`	Whether this parameter needs to be initialized (not loaded from a checkpoint).
init_fn	`Optional[Callable]`	`None`	Custom initialization function for parameters that need initialization.
dim	`int`	`0`	The dimension along which this parameter will be partitioned for tensor parallelism.
mapping	`Dict[HashableSlice, ParameterSlice]`	`{}`	Mapping from tensor slices to their source locations in weight files.

Supporting Types

HashableSlice (core.py L25-43)

A hashable wrapper around Python's slice object, enabling use as dictionary keys in the mapping field.

@dataclass
class HashableSlice:
    start: Optional[int] = None
    stop: Optional[int] = None
    step: Optional[int] = None

ParameterSlice (core.py L46-66)

Records where a specific slice of a parameter can be found in the original weight files.

@dataclass
class ParameterSlice:
    source: str          # Source file path
    shape: Tuple[int]    # Original shape of the full parameter
    slice_: HashableSlice  # Which slice of the source to read

Behavior

The initialize_parameter_meta function performs these steps:

Walk the module tree: Iterate over all named parameters in the model using model.named_parameters().
Detect tied parameters: Track seen parameter identities (via data_ptr() or object identity for meta tensors). If a parameter has been seen before, mark it as tied.
Attach ParameterMeta: For each unique parameter, create a new ParameterMeta instance and attach it as an attribute.
Mark tied pairs: When a tied parameter is found, mark both the original and the duplicate as is_tied=True.

Example Usage

from transformers import AutoConfig, AutoModelForCausalLM
from optimum.fx.parallelization.utils import MetaAwareMethodsPatcher, initialize_parameter_meta

config = AutoConfig.from_pretrained("meta-llama/Llama-2-7b-hf")

with MetaAwareMethodsPatcher():
    model = AutoModelForCausalLM.from_config(config)

# Attach metadata to all parameters
initialize_parameter_meta(model)

# Inspect metadata on a specific parameter
for name, param in model.named_parameters():
    meta = param.meta
    print(f"{name}: tied={meta.is_tied}, dim={meta.dim}")
    break

Page Connections

Double-click a node to navigate. Hold to expand connections.

Principle

Implementation

Heuristic

Environment