nemo_automodel.shared.tp_linear

Low-level linear graph-shaping helpers shared by tensor-parallel components.

Inductor’s async tensor-parallel pass (torch._inductor.config._micro_pipeline_tp) fuses collectives with matmuls by pattern-matching the reshape-mm-reshape graph that F.linear produces for 3-D input. The compile-safe torch.bmm path used by TPLinear/LinearLoRA never matches, so async-TP fusion silently fails to fire. These helpers detect async-TP tracing and emit the native linear graph in that mode only.

Module Contents

Functions

Name	Description
`_async_tp_linear`	Emit the native linear graph recognized by async-TP fusion.
`_is_async_tp_linear_enabled`	Return whether Dynamo is tracing with Inductor’s async-TP pass enabled.
`tp_linear_forward`	Dispatch a TP-safe linear between F.linear, async-TP shaping, and bmm/mm.

API

nemo_automodel.shared.tp_linear._async_tp_linear(
    x: torch.Tensor,
    weight: torch.Tensor,
    bias: torch.Tensor | None = None
) -> torch.Tensor

Emit the native linear graph recognized by async-TP fusion.

Parameters:

torch.Tensor

Input activations of shape [..., in_features]; any number of leading dimensions is accepted (typically [B, S, in_features] with B = batch, S = sequence). May be a DTensor with tensor-parallel placements.

weight

torch.Tensor

Weight of shape [out_features, in_features]; may be a DTensor sharded for colwise (Shard(0)) or rowwise (Shard(1)) tensor parallelism.

bias

torch.Tensor | NoneDefaults to None

Optional bias of shape [out_features].

Returns: torch.Tensor

Output of shape [..., out_features] with the same leading

nemo_automodel.shared.tp_linear._is_async_tp_linear_enabled() -> bool

Return whether Dynamo is tracing with Inductor’s async-TP pass enabled.

True only when both conditions hold: the caller is being traced by torch.compile and torch._inductor.config._micro_pipeline_tp is set (see enable_async_tensor_parallel in parallelizer.py). Always False in eager mode.

nemo_automodel.shared.tp_linear.tp_linear_forward(
    x: torch.Tensor,
    weight: torch.Tensor,
    bias: torch.Tensor | None,
    mm_for_2d_compile: bool
) -> torch.Tensor

Dispatch a TP-safe linear between F.linear, async-TP shaping, and bmm/mm.

Shared by TPLinear.forward and LinearLoRA.forward. Eager unsharded/last-dimension-sharded inputs use F.linear; async-TP tracing emits the fusable native linear graph via _async_tp_linear; every remaining compile path and dim-0/1-sharded DTensor input keeps the DTensor-safe matmul fallback (aten.view cannot flatten a sharded dimension).

Parameters:

torch.Tensor

Input activations of shape [B, S, in_features] (B = batch, S = sequence) or [N, in_features] (N = flattened tokens). May be a DTensor: a 3-D DTensor sharded on dim 0 or 1 (e.g. Shard(1) from sequence parallelism) always takes the torch.bmm path — the async-TP fast path assumes the input is replicated or sharded only on the feature dim, so fusion simply does not fire for such layers.

weight

torch.Tensor

Weight of shape [out_features, in_features]; may be a DTensor sharded for colwise (Shard(0)) or rowwise (Shard(1)) tensor parallelism.

bias

torch.Tensor | None

Optional bias of shape [out_features]; replicated if DTensor.

mm_for_2d_compile

bool

Numerics for 2-D x traced under torch.compile without async-TP: torch.mm plus explicit bias add when True (TPLinear), F.linear when False (LinearLoRA).

Returns: torch.Tensor

Output of shape [..., out_features] with the same leading