`bridge.models.conversion.fp8_export`#

Topology-neutral helpers for direct Transformer Engine FP8 parameter export.

Module Contents#

Classes#

FP8ExportLayout

Topology-neutral FP8 parameter layout used during checkpoint export.

Functions#

`detect_fp8_export_layout`	Inspect TE metadata without validating before the bridge’s PP collective.
`get_fp8_export_tensors`	Extract FP8 weight and inverse-scale tensors after collective layout validation.

Data#

logger

API#

bridge.models.conversion.fp8_export.logger#: ‘getLogger(…)’

class bridge.models.conversion.fp8_export.FP8ExportLayout#

Topology-neutral FP8 parameter layout used during checkpoint export.

format_name: str#: None

fp8_dtype: str | None#: None

block_shape: tuple[int | None, int | None]#: None

data_dtype: torch.dtype | None#: None

scale_dtype: torch.dtype | None#: None

scale_shape: tuple[int, ...]#: None

compact_scale_shape: tuple[int, int] | None#: None

with_gemm_swizzled_scales: bool#: None

validate() → None#: Validate format-specific export invariants.

bridge.models.conversion.fp8_export.detect_fp8_export_layout( local_weights: torch.Tensor, *, fp8_recipe: str | None, fp8_scale_inv_attr: str, ) → bridge.models.conversion.fp8_export.FP8ExportLayout | None#

Inspect TE metadata without validating before the bridge’s PP collective.

Parameters:

local_weights – Local parameter to inspect.
fp8_recipe – Configured FP8 recipe.
fp8_scale_inv_attr – Scale metadata key, optionally prefixed with an underscore.

Returns:

Layout for a supported FP8 parameter, or None when no usable scale metadata exists.

bridge.models.conversion.fp8_export.get_fp8_export_tensors( local_weights: torch.Tensor | None, *, fp8_scale_inv_attr: str, ) → tuple[torch.Tensor | None, torch.Tensor | None]#

Extract FP8 weight and inverse-scale tensors after collective layout validation.