diff --git a/megatron/core/distributed/param_and_grad_buffer.py b/megatron/core/distributed/param_and_grad_buffer.py index 33eb335c1fa..fe0a51b86b7 100644 --- a/megatron/core/distributed/param_and_grad_buffer.py +++ b/megatron/core/distributed/param_and_grad_buffer.py @@ -867,25 +867,63 @@ def _pad_start_of_param_if_needed(param_start_index: int) -> int: per_bucket_numel_unpadded = [] bucket_id = 0 - def _update_bucket_metadata(param_end_index: int) -> int: + def _update_bucket_metadata( + param_end_index: int, + bucket_start_index: int, + bucket_indices: list, + numel_unpadded_list: list, + ) -> int: """ - Record metadata for the bucket starting at bucket_start_index and ending with the - passed-in param_end_index. Returns the bucket's end_index. + Record metadata for a bucket. Returns the bucket's (padded) end_index. + + Args: + param_end_index: End index of the last param in this bucket (unpadded). + bucket_start_index: Start index of this bucket. + bucket_indices: List to append (start, end) bucket boundaries to. + numel_unpadded_list: List to append unpadded bucket numel to. + + Returns: + The bucket's end index, padded if using distributed optimizer. """ - nonlocal bucket_start_index, bucket_params, bucket_id - per_bucket_numel_unpadded.append(param_end_index - bucket_start_index) + numel_unpadded_list.append(param_end_index - bucket_start_index) bucket_end_index = _pad_end_of_bucket_if_needed(param_end_index) + bucket_indices.append((bucket_start_index, bucket_end_index)) + return bucket_end_index + + def _finalize_bucket_all_index_spaces( + param_end_index: int, + bucket_start_index: int, + nvfp4_packed_param_end_index: int = None, + nvfp4_packed_bucket_start_index: int = None, + ) -> tuple: + """ + Record metadata for the current bucket across both main and (if applicable) + NVFP4 packed index spaces. Also resets bucket_params and increments bucket_id. - # Record metadata of new bucket. - self.bucket_indices.append((bucket_start_index, bucket_end_index)) - bucket_start_index = bucket_end_index + Args: + param_end_index: End index of the last param in the bucket (full numel). + bucket_start_index: Start index of the bucket (full numel). + nvfp4_packed_param_end_index: End index in packed space (NVFP4 only). + nvfp4_packed_bucket_start_index: Bucket start in packed space (NVFP4 only). - # Prepare for next bucket. + Returns: + Tuple of (bucket_end_index, nvfp4_packed_bucket_end_index). + """ + nonlocal bucket_params, bucket_id + bucket_end_index = _update_bucket_metadata( + param_end_index, bucket_start_index, self.bucket_indices, per_bucket_numel_unpadded + ) + nvfp4_packed_bucket_end_index = None + if self.has_nvfp4_params: + nvfp4_packed_bucket_end_index = _update_bucket_metadata( + nvfp4_packed_param_end_index, + nvfp4_packed_bucket_start_index, + self.nvfp4_packed_bucket_indices, + nvfp4_packed_per_bucket_numel_unpadded, + ) bucket_params = set() bucket_id += 1 - - # Return the potentially padded bucket_end_index. - return bucket_end_index + return bucket_end_index, nvfp4_packed_bucket_end_index def _does_param_require_new_bucket(param): """ @@ -915,82 +953,110 @@ def _does_param_require_new_bucket(param): # Grad buffer: [g0, g1, g2, g3, ...] numel = N # # We therefore maintain two index maps: - # - param_index_map: offsets into the packed param buffer (numel // 2) - # - nvfp4_unpacked_param_index_map: offsets using full (unpacked) numel + # - param_index_map: offsets using full numel. + # - nvfp4_packed_param_index_map: offsets into the packed param buffer (numel // 2). # self.has_nvfp4_params = any(is_nvfp4tensor(p) for p in self.params) - self.nvfp4_unpacked_param_index_map = {} - grad_start_index = 0 if self.has_nvfp4_params else None - grad_bucket_start_index = 0 if self.has_nvfp4_params else None - grad_bucket_end_index = None + # Secondary (packed) index map, counters, and bucket tracking for NVFP4. + self.nvfp4_packed_param_index_map = {} if self.has_nvfp4_params else None + nvfp4_packed_param_start_index = 0 if self.has_nvfp4_params else None + nvfp4_packed_param_end_index = None + nvfp4_packed_bucket_start_index = 0 if self.has_nvfp4_params else None + self.nvfp4_packed_bucket_indices = [] if self.has_nvfp4_params else None + nvfp4_packed_per_bucket_numel_unpadded = [] if self.has_nvfp4_params else None for param, _ in params_with_names[::-1]: # Iterate through parameters in reverse order to roughly follow backprop order. - full_numel = param.data.nelement() - # NVFP4 params are packed (2 values per byte), so the param buffer uses - # half the logical numel. Non-NVFP4 params use the full numel for both. - if self.has_nvfp4_params and is_nvfp4tensor(param): - assert ( - full_numel % 2 == 0 - ), f"NVFP4 requires even numel for packing, got {full_numel}" - param_numel = full_numel // 2 - else: - param_numel = full_numel param_start_index = _pad_start_of_param_if_needed(param_start_index) + if self.has_nvfp4_params: + nvfp4_packed_param_start_index = _pad_start_of_param_if_needed( + nvfp4_packed_param_start_index + ) # Create bucket with collected parameters if current param needs its own bucket. if _does_param_require_new_bucket(param) and len(bucket_params) > 0: - # Ensure this param accounts for the new padding introduced at end of - # previous bucket. - param_start_index = _update_bucket_metadata(param_start_index) + # Finalize the current bucket and update start indices for the next bucket. + bucket_start_index, nvfp4_packed_bucket_start_index = ( + _finalize_bucket_all_index_spaces( + param_start_index, + bucket_start_index, + nvfp4_packed_param_start_index, + nvfp4_packed_bucket_start_index, + ) + ) + param_start_index = bucket_start_index + if self.has_nvfp4_params: + nvfp4_packed_param_start_index = nvfp4_packed_bucket_start_index + # Primary index computation: always uses full param numel. + param_numel = param.data.nelement() param_end_index = param_start_index + param_numel self.param_index_map[param] = (param_start_index, param_end_index, bucket_id) - bucket_params.add(param) - # For NVFP4, the grad buffer is sized at full numel (not packed), so we - # maintain a parallel index map using full_numel for every param. + # Secondary (packed) index computation for NVFP4. if self.has_nvfp4_params: - grad_start_index = _pad_start_of_param_if_needed(grad_start_index) - grad_end_index = grad_start_index + full_numel - self.nvfp4_unpacked_param_index_map[param] = ( - grad_start_index, - grad_end_index, + if is_nvfp4tensor(param): + assert ( + param_numel % 2 == 0 + ), f"NVFP4 requires even numel for packing, got {param_numel}" + # NVFP4 packs two FP4 values into one byte, so packed numel is half. + nvfp4_packed_param_end_index = nvfp4_packed_param_start_index + param_numel // 2 + else: + nvfp4_packed_param_end_index = nvfp4_packed_param_start_index + param_numel + self.nvfp4_packed_param_index_map[param] = ( + nvfp4_packed_param_start_index, + nvfp4_packed_param_end_index, bucket_id, ) - grad_start_index = grad_end_index + + bucket_params.add(param) # If we have enough elements already or the current param is part of the shared # embedding layer and needs a separate bucket, form a new bucket. if ( bucket_size is not None and (param_end_index - bucket_start_index) >= bucket_size ) or _does_param_require_new_bucket(param): - bucket_end_index = _update_bucket_metadata(param_end_index) - param_start_index = bucket_end_index + # Finalize the current bucket and update start indices for the next bucket. + bucket_start_index, nvfp4_packed_bucket_start_index = ( + _finalize_bucket_all_index_spaces( + param_end_index, + bucket_start_index, + nvfp4_packed_param_end_index, + nvfp4_packed_bucket_start_index, + ) + ) + param_start_index = bucket_start_index + if self.has_nvfp4_params: + nvfp4_packed_param_start_index = nvfp4_packed_bucket_start_index else: param_start_index = param_end_index + if self.has_nvfp4_params: + nvfp4_packed_param_start_index = nvfp4_packed_param_end_index # Add remaining params to a new bucket. if len(bucket_params) > 0: - bucket_end_index = _update_bucket_metadata(param_end_index) + _finalize_bucket_all_index_spaces( + param_end_index, + bucket_start_index, + nvfp4_packed_param_end_index, + nvfp4_packed_bucket_start_index, + ) # Next, create underlying storage for buffer (with numel elements that includes # padding as necessary). - self.numel = bucket_end_index + self.numel = self.bucket_indices[-1][1] self.numel_unpadded = sum(per_bucket_numel_unpadded) - - # For NVFP4, grad buffer needs full size (roughly 2x the packed param buffer). if self.has_nvfp4_params: - self.grad_numel = grad_start_index - if self.ddp_config.use_distributed_optimizer: - self.grad_numel = _pad(self.grad_numel, self.data_parallel_world_size) - else: - self.grad_numel = self.numel + self.nvfp4_packed_numel = self.nvfp4_packed_bucket_indices[-1][1] + self.nvfp4_packed_numel_unpadded = sum(nvfp4_packed_per_bucket_numel_unpadded) assert self.numel_unpadded <= self.numel if self.ddp_config.use_distributed_optimizer: assert self.numel % self.data_parallel_world_size == 0 + if self.has_nvfp4_params: + assert self.nvfp4_packed_numel_unpadded <= self.nvfp4_packed_numel + assert self.nvfp4_packed_numel % self.data_parallel_world_size == 0 else: assert self.numel == self.numel_unpadded @@ -1044,16 +1110,15 @@ def _does_param_require_new_bucket(param): else: # Only re-map param tensors if using distributed optimizer. if self.ddp_config.use_distributed_optimizer: + numel = self.nvfp4_packed_numel if self.has_nvfp4_params else self.numel self.param_data = torch.zeros( - self.numel, + numel, dtype=self.param_dtype, device=torch.cuda.current_device(), requires_grad=False, ) - # For NVFP4, grad buffer uses full size (grad_numel), - # param buffer uses packed size (numel) self.grad_data = torch.zeros( - self.grad_numel, + self.numel, dtype=self.grad_dtype, device=torch.cuda.current_device(), requires_grad=False, @@ -1064,12 +1129,46 @@ def _does_param_require_new_bucket(param): self.param_data_cpu = None # Finally, map param.data and param.main_grad fields to buffers. + def _create_bucket(bucket_id, bucket_params, bucket_params_with_extra_main_grads): + """ + Look up precomputed bucket indices and create a new bucket. + + Args: + bucket_id: ID of the bucket to create. + bucket_params: List of parameters in this bucket. + bucket_params_with_extra_main_grads: List of parameters with + extra FP32 main_grads. + + Returns: + A new _ParamAndGradBucket instance. + """ + bucket_start_index, bucket_end_index = self.bucket_indices[bucket_id] + if self.has_nvfp4_params: + nvfp4_packed_start_index, nvfp4_packed_end_index = self.nvfp4_packed_bucket_indices[ + bucket_id + ] + else: + nvfp4_packed_start_index, nvfp4_packed_end_index = None, None + return self._new_bucket( + bucket_params=bucket_params, + start_index=bucket_start_index, + end_index=bucket_end_index, + numel_unpadded=per_bucket_numel_unpadded[bucket_id], + bucket_id=bucket_id, + nvfp4_packed_start_index=nvfp4_packed_start_index, + nvfp4_packed_end_index=nvfp4_packed_end_index, + bucket_params_with_extra_main_grads=bucket_params_with_extra_main_grads, + ) + bucket_params = [] bucket_params_with_extra_main_grads = [] - bucket_start_index = 0 cur_bucket_id = 0 for param, param_name in params_with_names[::-1]: + # Get parameter indices computed in previous loop. param_start_index, param_end_index, bucket_id = self.param_index_map[param] + nvfp4_packed_param_start_index = None + if self.has_nvfp4_params: + nvfp4_packed_param_start_index, _, _ = self.nvfp4_packed_param_index_map[param] # For MXFP8 param: # we only need to map bf16 weights (layernorm, embedding, etc) to the buffer. if not self.ddp_config.reuse_grad_buf_for_mxfp8_param_ag or not is_mxfp8tensor(param): @@ -1082,17 +1181,31 @@ def _does_param_require_new_bucket(param): packed_shape = get_nvfp4_rowwise_packed_shape(param.data.shape) rowwise_bytes_view = self._get( - packed_shape, param_start_index, buffer_type=BufferType.PARAM + packed_shape, + nvfp4_packed_param_start_index, + buffer_type=BufferType.PARAM, ) modify_nvfp4_rowwise_storage(param, rowwise_bytes_view) elif is_float8tensor(param): new_param_data = self._get( - param.data.shape, param_start_index, buffer_type=BufferType.PARAM + param.data.shape, + ( + nvfp4_packed_param_start_index + if self.has_nvfp4_params + else param_start_index + ), + buffer_type=BufferType.PARAM, ) modify_underlying_storage(param, new_param_data) else: new_param_data = self._get( - param.data.shape, param_start_index, buffer_type=BufferType.PARAM + param.data.shape, + ( + nvfp4_packed_param_start_index + if self.has_nvfp4_params + else param_start_index + ), + buffer_type=BufferType.PARAM, ) old_param_data = param.data param.data = new_param_data @@ -1101,16 +1214,10 @@ def _does_param_require_new_bucket(param): param.data.detach().copy_(old_param_data) del old_param_data - # For NVFP4, use grad_index_map for main_grad (full numel offsets) - if self.has_nvfp4_params: - grad_start, grad_end, _ = self.nvfp4_unpacked_param_index_map[param] - param.main_grad = self._get( - param.data.shape, grad_start, buffer_type=BufferType.GRAD - ) - else: - param.main_grad = self._get( - param.data.shape, param_start_index, buffer_type=BufferType.GRAD - ) + # Grad buffer always uses full-numel offsets from param_index_map. + param.main_grad = self._get( + param.data.shape, param_start_index, buffer_type=BufferType.GRAD + ) # Create FP32 copy of .main_grads if necessary. promote_main_grads_to_higher_precision = False for param_name_pattern in ddp_config.param_name_patterns_for_fp32_local_accumulation: @@ -1135,24 +1242,11 @@ def _does_param_require_new_bucket(param): self.extra_main_grads.append(param.main_grad) if bucket_id != cur_bucket_id: - bucket_end_index = _pad_end_of_bucket_if_needed(param_start_index) - if self.has_nvfp4_params: - grad_bucket_end_index = _pad_end_of_bucket_if_needed(grad_start) self.buckets.append( - self._new_bucket( - bucket_params=bucket_params, - start_index=bucket_start_index, - end_index=bucket_end_index, - numel_unpadded=per_bucket_numel_unpadded[cur_bucket_id], - bucket_id=cur_bucket_id, - grad_start_index=grad_bucket_start_index, - grad_end_index=grad_bucket_end_index, - bucket_params_with_extra_main_grads=bucket_params_with_extra_main_grads, + _create_bucket( + cur_bucket_id, bucket_params, bucket_params_with_extra_main_grads ) ) - if self.has_nvfp4_params: - grad_bucket_start_index = grad_bucket_end_index - bucket_start_index = bucket_end_index bucket_params = [] bucket_params_with_extra_main_grads = [] assert cur_bucket_id + 1 == len(self.buckets) @@ -1165,20 +1259,8 @@ def _does_param_require_new_bucket(param): # Add remaining params to a new bucket. if len(bucket_params) > 0: - bucket_end_index = _pad_end_of_bucket_if_needed(param_end_index) - if self.has_nvfp4_params: - grad_bucket_end_index = self.grad_numel self.buckets.append( - self._new_bucket( - bucket_params=bucket_params, - start_index=bucket_start_index, - end_index=bucket_end_index, - numel_unpadded=per_bucket_numel_unpadded[cur_bucket_id], - bucket_id=cur_bucket_id, - grad_start_index=grad_bucket_start_index, - grad_end_index=grad_bucket_end_index, - bucket_params_with_extra_main_grads=bucket_params_with_extra_main_grads, - ) + _create_bucket(cur_bucket_id, bucket_params, bucket_params_with_extra_main_grads) ) # Log buckets for all PP stages. log_strs = [] @@ -1210,22 +1292,6 @@ def scale_gradients(self, scaling_factor: float) -> None: for grad in self.extra_main_grads: grad *= scaling_factor - def get_unpacked_index_map(self) -> Dict[torch.nn.Parameter, tuple[int, int, int]]: - """ - Return the index map using unpacked (full) numel for each parameter. - - For NVFP4 buffers, param_index_map uses packed numel (half the logical size), - so this returns nvfp4_unpacked_param_index_map which has full-numel indices instead. - For other buffers, packed and unpacked indices are identical, so param_index_map - is returned directly. - - The distributed optimizer uses this to determine which rank owns which portion - of each parameter's data. - """ - if self.has_nvfp4_params: - return self.nvfp4_unpacked_param_index_map - return self.param_index_map - def _get(self, shape: torch.Size, start_index: int, buffer_type: BufferType) -> torch.Tensor: """ Return a tensor with the input `shape` as a view into the 1-D data starting at @@ -1233,11 +1299,12 @@ def _get(self, shape: torch.Size, start_index: int, buffer_type: BufferType) -> """ end_index = start_index + shape.numel() if buffer_type == BufferType.PARAM: - assert end_index <= self.numel, "Requested tensor is out of param buffer range" + numel = self.nvfp4_packed_numel if self.has_nvfp4_params else self.numel + assert end_index <= numel, "Requested tensor is out of param buffer range" assert self.param_data is not None buffer_tensor = self.param_data[start_index:end_index] elif buffer_type == BufferType.GRAD: - assert end_index <= self.grad_numel, "Requested tensor is out of grad buffer range" + assert end_index <= self.numel, "Requested tensor is out of grad buffer range" buffer_tensor = self.grad_data[start_index:end_index] else: raise Exception("Illegal buffer type provided to GradBuffer._get() function") @@ -1252,14 +1319,15 @@ def _new_bucket( numel_unpadded: int, bucket_id: int, bucket_params_with_extra_main_grads: List[torch.Tensor], - grad_start_index: int = None, - grad_end_index: int = None, + nvfp4_packed_start_index: int = None, + nvfp4_packed_end_index: int = None, ) -> _ParamAndGradBucket: """ Helper function that creates a new bucket. Also updates param->bucket mapping. - For NVFP4 buffers, grad_start_index and grad_end_index are provided separately - because grad buffer uses full numel while param buffer uses packed numel. + For NVFP4 buffers, nvfp4_packed_start_index and nvfp4_packed_end_index + are provided separately because the param buffer uses packed numel while + the grad buffer uses full numel. """ # Assert that indices are correctly padded (if needed), and that bucket @@ -1269,32 +1337,36 @@ def _new_bucket( assert start_index % self.data_parallel_world_size == 0 assert end_index % self.data_parallel_world_size == 0 assert (start_index, end_index) == self.bucket_indices[bucket_id] + if nvfp4_packed_start_index is not None: + assert ( + nvfp4_packed_start_index, + nvfp4_packed_end_index, + ) == self.nvfp4_packed_bucket_indices[bucket_id] # Get appropriate view into global _ParamAndGradBuffer. + # For NVFP4, param buffer uses packed offsets; otherwise same as start/end. bucketed_param_data = None if self.param_data is not None: - bucketed_param_data = self._get( - torch.Size([end_index - start_index]), start_index, buffer_type=BufferType.PARAM - ) - # For NVFP4, use separate grad buffer offsets - if grad_start_index is not None and grad_end_index is not None: - bucketed_grad_data = self._get( - torch.Size([grad_end_index - grad_start_index]), - grad_start_index, - buffer_type=BufferType.GRAD, - ) - else: - bucketed_grad_data = self._get( - torch.Size([end_index - start_index]), start_index, buffer_type=BufferType.GRAD - ) - # For NVFP4, use grad buffer offset for bucket.offset since distrib_optimizer - # uses it for grad buffer operations. For non-NVFP4, param and grad offsets are same. - bucket_offset = grad_start_index if grad_start_index is not None else start_index + if nvfp4_packed_start_index is not None: + assert nvfp4_packed_end_index is not None + bucketed_param_data = self._get( + torch.Size([nvfp4_packed_end_index - nvfp4_packed_start_index]), + nvfp4_packed_start_index, + buffer_type=BufferType.PARAM, + ) + else: + bucketed_param_data = self._get( + torch.Size([end_index - start_index]), start_index, buffer_type=BufferType.PARAM + ) + # Grad buffer always uses full-numel offsets. + bucketed_grad_data = self._get( + torch.Size([end_index - start_index]), start_index, buffer_type=BufferType.GRAD + ) bucket = _ParamAndGradBucket( params=bucket_params, param_data=bucketed_param_data, grad_data=bucketed_grad_data, - offset=bucket_offset, + offset=start_index, numel_unpadded=numel_unpadded, gradient_scaling_factor=self.gradient_scaling_factor, bucket_id=bucket_id, diff --git a/megatron/core/optimizer/distrib_optimizer.py b/megatron/core/optimizer/distrib_optimizer.py index a45eceefe35..b4d52e6b56b 100644 --- a/megatron/core/optimizer/distrib_optimizer.py +++ b/megatron/core/optimizer/distrib_optimizer.py @@ -218,10 +218,8 @@ def _build_model_gbuf_range(cls, param_and_grad_buffer: _ParamAndGradBuffer, buc gbuf_world_range = gbuf_world_all_ranges[data_parallel_rank] # Get each param's ranges. - # Use get_unpacked_index_map() which returns full-numel indices for NVFP4 params - # (from nvfp4_unpacked_param_index_map) and normal indices for other params. param_range_map = cls._build_model_gbuf_param_range_map( - param_and_grad_buffer.get_unpacked_index_map(), gbuf_world_range, bucket.offset + param_and_grad_buffer.param_index_map, gbuf_world_range, bucket.offset ) # Group into dict. @@ -1504,6 +1502,9 @@ def sharded_param_state_fully_reshardable( for dtype, gbuf_range_map_for_all_buckets in gbuf_range_maps.items(): world_tensors = dp_zero_state_dict[gbuf_idx][dtype] world_tensor_keys = world_tensors.keys() + # Note: for NVFP4, param_index_map uses unpacked (full numel) + # offsets, which is correct here since optimizer states + # (fp32_param, exp_avg, exp_avg_sq) are in unpacked space. for model_param, ( param_world_start, param_world_end, diff --git a/tests/unit_tests/dist_checkpointing/test_optimizer.py b/tests/unit_tests/dist_checkpointing/test_optimizer.py index 97928bd5954..1db844dd9a8 100644 --- a/tests/unit_tests/dist_checkpointing/test_optimizer.py +++ b/tests/unit_tests/dist_checkpointing/test_optimizer.py @@ -591,7 +591,7 @@ def test_bucket_space_optimizer_save_load( # Init model and optimizer with "src" bucket padding with patch('megatron.core.distributed.param_and_grad_buffer.math.lcm') as lcm_mock: lcm_mock.return_value = src_bucket_pad_divisor - assert len(lcm_mock.mock_calls) == 0 + model_A, optimizer_A = setup_model_and_optimizer( seed=2, tp=src_tp_pp[0], @@ -600,7 +600,6 @@ def test_bucket_space_optimizer_save_load( dist_opt=True, initialize_fn=initialize_pp_agnostic_model, ) - assert len(lcm_mock.mock_calls) > 1 metadata = {'distrib_optim_sharding_type': 'dp_reshardable'} @@ -618,7 +617,7 @@ def test_bucket_space_optimizer_save_load( # Init model and optimizer with "dest" bucket padding with patch('megatron.core.distributed.param_and_grad_buffer.math.lcm') as lcm_mock: lcm_mock.return_value = dest_bucket_pad_divisor - assert len(lcm_mock.mock_calls) == 0 + model_B, optimizer_B = setup_model_and_optimizer( seed=3, tp=dest_tp_pp[0], @@ -627,7 +626,6 @@ def test_bucket_space_optimizer_save_load( dist_opt=True, initialize_fn=initialize_pp_agnostic_model, ) - assert len(lcm_mock.mock_calls) > 1 model_sharded_sd = model_B[0].sharded_state_dict() load_sharded_state_dict = optimizer_B.sharded_state_dict( diff --git a/tests/unit_tests/distributed/test_param_and_grad_buffer.py b/tests/unit_tests/distributed/test_param_and_grad_buffer.py index 720c5a2d1af..223383bba64 100644 --- a/tests/unit_tests/distributed/test_param_and_grad_buffer.py +++ b/tests/unit_tests/distributed/test_param_and_grad_buffer.py @@ -10,7 +10,7 @@ from megatron.core import parallel_state from megatron.core.distributed import DistributedDataParallel, DistributedDataParallelConfig -from megatron.core.distributed.param_and_grad_buffer import partition_buckets +from megatron.core.distributed.param_and_grad_buffer import _ParamAndGradBuffer, partition_buckets from megatron.core.transformer import TransformerConfig from tests.unit_tests.test_utilities import TestModel, Utils @@ -669,3 +669,218 @@ def test_grad_sync_copies_to_and_from_comm_buffer(self): ) Utils.destroy_model_parallel() + + +class TestNVFP4IndexMaps: + """Tests for NVFP4 dual index map (param_index_map and nvfp4_packed_param_index_map). + + These tests mock NVFP4 functions and CUDA so they run on CPU without GPUs. + The mocking replaces is_nvfp4tensor (to treat regular bf16 params as NVFP4), + get_nvfp4_rowwise_packed_shape (to halve last dim), modify_nvfp4_rowwise_storage + (no-op), and torch.cuda.current_device (to allocate on CPU). + """ + + @staticmethod + def _make_buffer( + param_shapes, + nvfp4_param_indices=None, + use_distributed_optimizer=False, + bucket_size=None, + dp_world_size=1, + ): + """Create a _ParamAndGradBuffer with some params mocked as NVFP4. + + Args: + param_shapes: List of (name, shape) tuples for each parameter. + nvfp4_param_indices: Set of indices into param_shapes to treat as NVFP4. + use_distributed_optimizer: Whether to use distributed optimizer. + bucket_size: Bucket size for splitting. + dp_world_size: Simulated data parallel world size. + + Returns: + (buffer, params) where params is the ordered list of nn.Parameters. + """ + params = [] + params_with_names = [] + param_to_name = {} + for name, shape in param_shapes: + param = torch.nn.Parameter(torch.randn(shape, dtype=torch.bfloat16)) + params.append(param) + params_with_names.append((param, name)) + param_to_name[param] = name + + if nvfp4_param_indices is None: + nvfp4_param_indices = set() + nvfp4_params = {params[i] for i in nvfp4_param_indices} + has_nvfp4 = len(nvfp4_params) > 0 + + def mock_is_nvfp4(t): + return any(t is p for p in nvfp4_params) + + def mock_packed_shape(shape): + packed = list(shape) + packed[-1] = packed[-1] // 2 + return torch.Size(packed) + + mock_dp_group = mock.MagicMock() + mock_dp_group.size.return_value = dp_world_size + mock_pg = mock.MagicMock() + + ddp_config = DistributedDataParallelConfig( + use_distributed_optimizer=use_distributed_optimizer, + overlap_grad_reduce=False, + bucket_size=bucket_size, + average_in_collective=False, + ) + + with ( + mock.patch( + 'megatron.core.distributed.param_and_grad_buffer.is_nvfp4tensor', + side_effect=mock_is_nvfp4, + ), + mock.patch( + 'megatron.core.distributed.param_and_grad_buffer.get_nvfp4_rowwise_packed_shape', + side_effect=mock_packed_shape, + ), + mock.patch('megatron.core.fp4_utils.modify_nvfp4_rowwise_storage'), + mock.patch('torch.cuda.current_device', return_value='cpu'), + mock.patch( + 'megatron.core.distributed.param_and_grad_buffer.log_on_each_pipeline_stage' + ), + ): + buffer = _ParamAndGradBuffer( + ddp_config=ddp_config, + param_dtype=torch.uint8 if has_nvfp4 else torch.bfloat16, + grad_dtype=torch.bfloat16, + params_with_names=params_with_names, + data_parallel_group=mock_dp_group, + bucket_size=bucket_size, + param_to_name=param_to_name, + gradient_scaling_factor=1.0, + param_indices=list(range(len(params))), + nccl_ub=False, + pg_collection=mock_pg, + ) + + return buffer, params + + def test_exact_index_values_no_padding(self): + """Verify exact index map values for a simple case without distributed optimizer.""" + param_shapes = [('layer0.weight', (100, 100)), ('layer1.weight', (100, 100))] + buffer, params = self._make_buffer(param_shapes, nvfp4_param_indices={0, 1}) + + # Buffer processes params in reverse order: params[1] first, params[0] second. + assert buffer.param_index_map[params[1]] == (0, 10000, 0) + assert buffer.param_index_map[params[0]] == (10000, 20000, 0) + assert buffer.nvfp4_packed_param_index_map[params[1]] == (0, 5000, 0) + assert buffer.nvfp4_packed_param_index_map[params[0]] == (5000, 10000, 0) + + assert buffer.numel == 20000 + assert buffer.nvfp4_packed_numel == 10000 + + def test_non_nvfp4_exact_values_match_original(self): + """Non-NVFP4 param_index_map values should be identical to original behavior.""" + param_shapes = [('layer0.weight', (100, 100)), ('layer1.weight', (100, 100))] + buffer, params = self._make_buffer(param_shapes) + + # Without NVFP4 or distributed optimizer, no padding: offsets are contiguous. + assert buffer.param_index_map[params[1]] == (0, 10000, 0) + assert buffer.param_index_map[params[0]] == (10000, 20000, 0) + assert buffer.numel == 20000 + + def test_nvfp4_multi_bucket_param_to_index(self): + """param_to_index in each bucket should be relative to that bucket's full-numel offset.""" + param_shapes = [ + ('layer0.weight', (100, 100)), + ('layer1.weight', (100, 100)), + ('layer2.weight', (100, 100)), + ('layer3.weight', (100, 100)), + ] + # bucket_size=15000: each param is 10000 full numel, so 2 params per bucket. + buffer, params = self._make_buffer( + param_shapes, nvfp4_param_indices={0, 1, 2, 3}, bucket_size=15000 + ) + + assert len(buffer.buckets) == 2 + for bucket in buffer.buckets: + for param in bucket.params_list: + global_start, global_end, _ = buffer.param_index_map[param] + local_start, local_end = bucket.param_to_index[param] + assert local_start == global_start - bucket.offset + assert local_end == global_end - bucket.offset + assert local_end - local_start == param.data.nelement() + + @pytest.mark.parametrize("dp_world_size", [1, 2, 4, 8]) + def test_nvfp4_with_distributed_optimizer(self, dp_world_size): + """With distributed optimizer, both packed and unpacked indices should be padded.""" + param_shapes = [('layer0.weight', (98, 101)), ('layer1.weight', (98, 101))] + buffer, params = self._make_buffer( + param_shapes, + nvfp4_param_indices={0, 1}, + use_distributed_optimizer=True, + dp_world_size=dp_world_size, + ) + + # Param starts should be 64-aligned in both maps. + for param in params: + start, end, _ = buffer.param_index_map[param] + assert start % 64 == 0, f"Unpacked start {start} should be 64-aligned" + assert end - start == param.data.nelement() + + for param in params: + start, end, _ = buffer.nvfp4_packed_param_index_map[param] + assert start % 64 == 0, f"Packed start {start} should be 64-aligned" + assert end - start == param.data.nelement() // 2 + + # Buffer numel should be divisible by dp_world_size. + assert buffer.numel % dp_world_size == 0 + assert buffer.nvfp4_packed_numel % dp_world_size == 0 + + def test_nvfp4_mixed_params(self): + """Test buffer with a mix of NVFP4 and non-NVFP4 params.""" + param_shapes = [ + ('linear.weight', (100, 100)), # NVFP4. + ('layernorm.weight', (100,)), # Non-NVFP4 (bf16). + ] + buffer, params = self._make_buffer(param_shapes, nvfp4_param_indices={0}) + + # Non-NVFP4 param should have same span in both maps. + packed_start, packed_end, _ = buffer.nvfp4_packed_param_index_map[params[1]] + unpacked_start, unpacked_end, _ = buffer.param_index_map[params[1]] + assert packed_end - packed_start == unpacked_end - unpacked_start == 100 + + # NVFP4 param should have half the span in packed map. + packed_start, packed_end, _ = buffer.nvfp4_packed_param_index_map[params[0]] + unpacked_start, unpacked_end, _ = buffer.param_index_map[params[0]] + assert packed_end - packed_start == 5000 # numel // 2. + assert unpacked_end - unpacked_start == 10000 # Full numel. + + def test_nvfp4_varied_param_sizes(self): + """Test with different param sizes to verify offsets accumulate correctly.""" + param_shapes = [('small.weight', (10, 20)), ('large.weight', (100, 200))] + buffer, params = self._make_buffer(param_shapes, nvfp4_param_indices={0, 1}) + + # Reversed order: large (params[1]) processed first, then small (params[0]). + large_packed_start = 0 + large_packed_end = 100 * 200 // 2 # 10000. + small_packed_start = large_packed_end + small_packed_end = small_packed_start + 10 * 20 // 2 # 10100. + + assert buffer.nvfp4_packed_param_index_map[params[1]] == ( + large_packed_start, + large_packed_end, + 0, + ) + assert buffer.nvfp4_packed_param_index_map[params[0]] == ( + small_packed_start, + small_packed_end, + 0, + ) + + large_unpacked_start = 0 + large_unpacked_end = 100 * 200 # 20000. + small_unpacked_start = large_unpacked_end + small_unpacked_end = small_unpacked_start + 10 * 20 # 20200. + + assert buffer.param_index_map[params[1]] == (large_unpacked_start, large_unpacked_end, 0) + assert buffer.param_index_map[params[0]] == (small_unpacked_start, small_unpacked_end, 0)