diff --git a/tests/pytorch/test_grouped_tensor.py b/tests/pytorch/test_grouped_tensor.py index ac8c493290..5e1bb40c69 100644 --- a/tests/pytorch/test_grouped_tensor.py +++ b/tests/pytorch/test_grouped_tensor.py @@ -4,8 +4,10 @@ """Tests for GroupedTensor class""" -from typing import List, Optional, Tuple import os +from types import SimpleNamespace +from typing import List, Optional, Tuple + import pytest import torch import transformer_engine.pytorch as te @@ -18,8 +20,7 @@ MXFP8Quantizer, NVFP4Quantizer, ) -from transformer_engine.pytorch.constants import TE_DType_To_Torch -from transformer_engine.pytorch.utils import is_non_tn_fp8_gemm_supported +from transformer_engine.pytorch.utils import is_non_tn_fp8_gemm_supported, mark_grouped_tensor import transformer_engine_torch as tex # Import test utilities @@ -48,6 +49,45 @@ ) ) + +def test_mark_grouped_tensor_supports_plain_tensor(): + tensor = torch.empty(16) + + mark_grouped_tensor(tensor) + + assert tensor.grouped_tensor_scale_inv is False + + +def test_mark_grouped_tensor_supports_unquantized_rowwise_storage(): + rowwise_data = torch.empty(16) + grouped_tensor = SimpleNamespace( + rowwise_data=rowwise_data, + columnwise_data=None, + columnwise_scale_inv=None, + ) + + mark_grouped_tensor(grouped_tensor) + + assert rowwise_data.grouped_tensor_scale_inv is False + + +def test_mark_grouped_tensor_marks_quantized_columnwise_storage(): + rowwise_data = torch.empty(16) + columnwise_data = torch.empty(16) + columnwise_scale_inv = torch.empty(4) + grouped_tensor = SimpleNamespace( + rowwise_data=rowwise_data, + columnwise_data=columnwise_data, + columnwise_scale_inv=columnwise_scale_inv, + ) + + mark_grouped_tensor(grouped_tensor) + + assert not hasattr(rowwise_data, "grouped_tensor_scale_inv") + assert columnwise_data.grouped_tensor_scale_inv is False + assert columnwise_scale_inv.grouped_tensor_scale_inv is True + + _quantization_params = [ pytest.param( "fp8_delayed_scaling", diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 612a430966..3ea05b8ffe 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -40,6 +40,7 @@ clear_tensor_data, get_device_compute_capability, init_method_constant, + mark_grouped_tensor, requires_grad, resolve_grouped_linear_single_param_flags, get_nvtx_range_context, @@ -553,6 +554,10 @@ def _forward_grouped_tensor( if not inp.requires_grad: weights_to_save = [None] * len(weights_to_save) + # Megatron-LM paged stashing uses this marker to identify the dynamic activation + # buffers among the tensors saved by the GroupedLinear autograd function. The + # operation-fuser grouped MLP applies the same marker to its saved activations. + mark_grouped_tensor(input_to_save) tensors_to_save, tensor_objects = prepare_for_saving( input_to_save, *weights_to_save, diff --git a/transformer_engine/pytorch/utils.py b/transformer_engine/pytorch/utils.py index 39162f8311..eb64412790 100644 --- a/transformer_engine/pytorch/utils.py +++ b/transformer_engine/pytorch/utils.py @@ -311,21 +311,39 @@ def mark_grouped_tensor(*tensors: List[Any]): Megatron-LM to detect which tensors are dynamic (varying shapes) and remove the padding before doing the `save_for_backward` to save memory. - Note: Only columnwise data is saved for backward.""" + + Plain tensors are saved directly. Grouped tensors are decomposed by + `prepare_for_saving`: unquantized BF16/FP16 activations save their + rowwise data, while quantized activations save their columnwise data + and scale metadata. + """ for tensor in tensors: if tensor is None: continue - if hasattr(tensor, "columnwise_data"): - assert ( - tensor.columnwise_data is not None - ), "Columnwise data is not set for grouped tensor" + + if not hasattr(tensor, "columnwise_data"): + # Plain tensor, e.g. a fused-MLP activation input. + setattr(tensor, "grouped_tensor_scale_inv", False) + continue + + # Grouped tensor: mark the underlying tensors that `prepare_for_saving` + # will pass to `save_for_backward`, rather than the storage wrapper. + if tensor.columnwise_data is None: + # Unquantized BF16/FP16 grouped tensor. + saved_activation = tensor.rowwise_data + saved_scale_inv = None + else: + # Quantized grouped tensor saved in the representation used by wgrad. + saved_activation = tensor.columnwise_data + saved_scale_inv = tensor.columnwise_scale_inv assert ( - tensor.columnwise_scale_inv is not None + saved_scale_inv is not None ), "Columnwise scale inverse is not set for grouped tensor" - setattr(tensor.columnwise_data, "grouped_tensor_scale_inv", False) - setattr(tensor.columnwise_scale_inv, "grouped_tensor_scale_inv", True) - else: - setattr(tensor, "grouped_tensor_scale_inv", False) + + assert saved_activation is not None, "Grouped tensor has no activation data" + setattr(saved_activation, "grouped_tensor_scale_inv", False) + if saved_scale_inv is not None: + setattr(saved_scale_inv, "grouped_tensor_scale_inv", True) def split_tensor_along_dim(