runtime_hidden_width

Differentiable full-envelope execution for a nested residual width.

Functions

hidden_width_layer_context	Execute one decoder block as a prefix-width child in a static PP envelope.
hidden_width_module_context	Apply descriptor-owned width rules to one module tree in a static envelope.
retained_hidden_prefix

hidden_width_layer_context(layer, *, canonical_layer_name, spec, width)

Execute one decoder block as a prefix-width child in a static PP envelope.

Residual inputs and outputs keep the full PP shape, while descriptor-owned input columns and output rows outside the active prefix are differentiably zeroed. RMSNorm is corrected to use the active-width denominator exactly. This preserves TP/FSDP module forwards and their collectives.

Parameters:

• layer (torch.nn.Module)

• canonical_layer_name (str)

• spec (EmbeddingPruningSpec)

• width (int)

hidden_width_module_context(module, *, canonical_module_name, spec, width, mask_boundary_input=False)

Apply descriptor-owned width rules to one module tree in a static envelope.

Parameters:

• module (torch.nn.Module)

• canonical_module_name (str)

• spec (EmbeddingPruningSpec)

• width (int)

• mask_boundary_input (bool)

retained_hidden_prefix(value, width)

Parameters:

• value (Tensor)

• width (int)

Return type:

Tensor