core.fp8_utils

Utility functions related to FP8 that are used throughout Megatron core

Module Contents

Functions

_unwrap_parameter_data	Return underlying tensor data when PyTorch wraps a tensor subclass as a Parameter.
_is_instance_or_param_data	Check a tensor subclass, including when wrapped by torch.nn.Parameter.
is_float8tensor	Check if a tensor is a Transformer Engine Float8Tensor.
is_mxfp8tensor	Check if a tensor is a Transformer Engine MXFP8Tensor
is_grouped_tensor	Check if a tensor is a Transformer Engine GroupedTensor.
is_grouped_tensor_with_quantized_storage	Check if a Transformer Engine GroupedTensor owns quantized primary storage.
_get_grouped_quantized_recipe	Return TE recipe for grouped quantized storage, or None if unavailable.
is_grouped_mxfp8tensor	Check if a TE GroupedTensor stores MXFP8 member tensors.
get_grouped_quantized_members	Return cached per-member views for a grouped quantized tensor.
copy_tensor_to_quantized_param	Copy high-precision values into TE quantized parameter storage.
copy_tensors_to_quantized_params	List form of :func:copy_tensor_to_quantized_param, for a whole bucket of params.
modify_grouped_tensor_rowwise_storage	Replace a high-precision Transformer Engine GroupedTensor’s rowwise storage.
dequantize_fp8_tensor	Dequantize a fp8 tensor to a higher precision tensor.
_resolve_callable_from_python_import_path	Resolve a Python import path like ‘pkg.mod.func’ to a callable.
_get_custom_recipe
get_fp8_align_size	Get the alignment size required for fp8 GEMM.
is_column_parallel_linear	Returns whether the given module is a ColumnParallelLinear layer.
is_row_parallel_linear	Returns whether the given module is a RowParallelLinear layer.
modify_underlying_storage	Replace the underlying raw data of a tensor with new data.
quantize_param_shard	Cast shard fp32 main params to fp8 model params.
correct_amax_history_if_needed	Correct the amax history of fp8 tensors when it’s necessary (i.e., in TE1.x).
post_all_gather_processing	Post-processing after all-gather for weights in distributed optimizer.
is_first_last_bf16_layer	Check if the layer is in bf16.
is_mxfp8_output_proj_active	Return True when the LM-head output projection should run under MXFP8.

Data

HAVE_TE
HAVE_TE_FP8_TENSOR_CLASS

API

core.fp8_utils.HAVE_TE

False

core.fp8_utils.HAVE_TE_FP8_TENSOR_CLASS

False

core.fp8_utils._unwrap_parameter_data(tensor: torch.Tensor) → torch.Tensor

Return underlying tensor data when PyTorch wraps a tensor subclass as a Parameter.

core.fp8_utils._is_instance_or_param_data(

tensor: torch.Tensor,

tensor_class: type,

) → bool

Check a tensor subclass, including when wrapped by torch.nn.Parameter.

core.fp8_utils.is_float8tensor(tensor: torch.Tensor) → bool

Check if a tensor is a Transformer Engine Float8Tensor.

Note that in TE2.x, in order to support more recipes, the design of the fp8 tensor class has changed. Now Float8Tensor is only used for current scaling and delayed scaling. And mxfp8 and blockwise scaling have their own fp8 tensor classes. These different fp8 tensor classes are both inherited from QuantizedTensor. So, for TE1.x, FP8_TENSOR_CLASS is Float8Tensor, and for TE2.x, FP8_TENSOR_CLASS is QuantizedTensor.

core.fp8_utils.is_mxfp8tensor(tensor: torch.Tensor) → bool

Check if a tensor is a Transformer Engine MXFP8Tensor

core.fp8_utils.is_grouped_tensor(tensor: torch.Tensor) → bool

Check if a tensor is a Transformer Engine GroupedTensor.

core.fp8_utils.is_grouped_tensor_with_quantized_storage(tensor: torch.Tensor) → bool

Check if a Transformer Engine GroupedTensor owns quantized primary storage.

core.fp8_utils._get_grouped_quantized_recipe(tensor: torch.Tensor)

Return TE recipe for grouped quantized storage, or None if unavailable.

core.fp8_utils.is_grouped_mxfp8tensor(tensor: torch.Tensor) → bool

Check if a TE GroupedTensor stores MXFP8 member tensors.

core.fp8_utils.get_grouped_quantized_members(

tensor: torch.Tensor,

*,

create_if_missing: bool = False,

) → List[torch.Tensor]

Return cached per-member views for a grouped quantized tensor.

core.fp8_utils.copy_tensor_to_quantized_param(

param: torch.Tensor,

src: torch.Tensor,

) → None

Copy high-precision values into TE quantized parameter storage.

core.fp8_utils.copy_tensors_to_quantized_params(

params: List[torch.Tensor],

srcs: List[torch.Tensor],

) → None

List form of :func:copy_tensor_to_quantized_param, for a whole bucket of params.

Same values, minus the per-param copy_ and tensor-subclass dispatch: the quantizer is resolved up front and called directly. Cast kernels are unchanged, one per param. Worth it because those casts are small and issuing them is expensive, and under –reuse-grad-buf-for-mxfp8-param-ag they run inside the forward pass.

Parameters:

• params – quantized model params to write into.

• srcs – high-precision source values, one per param, in the same order.

core.fp8_utils.modify_grouped_tensor_rowwise_storage(

tensor: torch.Tensor,

new_storage: torch.Tensor,

) → None

Replace a high-precision Transformer Engine GroupedTensor’s rowwise storage.

core.fp8_utils.dequantize_fp8_tensor(fp8_tensor: torch.Tensor) → torch.Tensor

Dequantize a fp8 tensor to a higher precision tensor.

core.fp8_utils._resolve_callable_from_python_import_path(dotted_path: str)

Resolve a Python import path like ‘pkg.mod.func’ to a callable.

Raises ValueError with clear message on failure.

core.fp8_utils._get_custom_recipe(

quantizer_factory_python_path: str,

) → Union[megatron.core.enums.Fp8Recipe, megatron.core.enums.Fp4Recipe]

core.fp8_utils.get_fp8_align_size(fp8_recipe: megatron.core.enums.Fp8Recipe) → int

Get the alignment size required for fp8 GEMM.

core.fp8_utils.is_column_parallel_linear(module)

Returns whether the given module is a ColumnParallelLinear layer.

core.fp8_utils.is_row_parallel_linear(module)

Returns whether the given module is a RowParallelLinear layer.

core.fp8_utils.modify_underlying_storage(

tensor: torch.Tensor,

new_raw_data: torch.Tensor,

)

Replace the underlying raw data of a tensor with new data.

core.fp8_utils.quantize_param_shard(

model_params,

main_params,

start_offsets,

data_parallel_group,

fsdp_shard_model_params=None,

)

Cast shard fp32 main params to fp8 model params.

core.fp8_utils.correct_amax_history_if_needed(model: List[torch.nn.Module])

Correct the amax history of fp8 tensors when it’s necessary (i.e., in TE1.x).

core.fp8_utils.post_all_gather_processing(model_params)

Post-processing after all-gather for weights in distributed optimizer.

• tensorwise: may need to create a transposed view to match backend GEMM.

• blockwise: create column-wise storage.

core.fp8_utils.is_first_last_bf16_layer(

config: megatron.core.transformer.transformer_config.TransformerConfig,

layer_no: int,

)

Check if the layer is in bf16.

core.fp8_utils.is_mxfp8_output_proj_active(config) → bool

Return True when the LM-head output projection should run under MXFP8.

Active when fp8_output_proj=True, fp8=True, fp8_recipe='mxfp8', and Transformer Engine is installed.