`bridge.models.common.heads`#

Module Contents#

Classes#

LinearForLastLayer

Replicated projection head compatible with Megatron output-layer calls.

Functions#

`create_value_head_hook`	Create a pre-wrap hook that replaces the final pipeline-stage output head.
`_create_last_layer_hook`
`ensure_model_list`	Normalize a single model chunk or a model-chunk list to a list.

Data#

`ModelList`
`ModelHook`

API#

bridge.models.common.heads.ModelList#: None

bridge.models.common.heads.ModelHook#: None

class bridge.models.common.heads.LinearForLastLayer( input_size: int, output_size: int, sequence_parallel: bool, bias: bool = False, dropout: float = 0.0, *, output_in_fp32: bool = True, tp_group: torch.distributed.ProcessGroup | None = None, init_method: collections.abc.Callable[[torch.Tensor], torch.Tensor | None] | None = None, perform_initialization: bool = True, )#

Bases: torch.nn.Linear

Replicated projection head compatible with Megatron output-layer calls.

Initialization

Initialize a replicated final projection.

Parameters:

input_size – Hidden dimension of the transformer output.
output_size – Output dimension of the projection head.
sequence_parallel – Whether to gather sequence-parallel activations.
bias – Whether to add a trainable bias.
dropout – Dropout probability applied before the projection.
output_in_fp32 – Whether to cast the projection output to FP32.
tp_group – Tensor-parallel process group used for sequence gathering.
init_method – Optional weight initializer for fresh-model construction.
perform_initialization – Whether to initialize the head parameters.

reset_parameters() → None#: Apply the configured initializer, including after meta-device materialization.

forward( input_: torch.Tensor, weight: torch.Tensor | None = None, runtime_gather_output: bool | None = None, ) → tuple[torch.Tensor, None]#: Run the final projection and return Megatron-style (output, bias).

bridge.models.common.heads.create_value_head_hook( hidden_size: int, sequence_parallel: bool, output_size: int = 1, ) → bridge.models.common.heads.ModelHook#

Create a pre-wrap hook that replaces the final pipeline-stage output head.