Make param_index_map always use unpacked (full numel) offsets

[deepakn94]

For NVFP4 buffers, param_index_map previously stored packed offsets (numel // 2). This was inconsistent with non-NVFP4 buffers and required callers to use get_unpacked_index_map() to get full-numel offsets.

Now param_index_map always stores unpacked offsets regardless of buffer type. The packed offsets are stored in nvfp4_packed_param_index_map, which is the NVFP4-specific concern. This also removes the now-trivial get_unpacked_index_map() method.

[github-actions]

This PR has been automatically converted to draft because all PRs must start as drafts.

When you are ready for review, click Ready for Review to begin the review process. This will:

1. Add the oncall reviewer (optional reviewer)
2. Add required review teams based on your changes

See the contribution guide for more details.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[deepakn94]

Note to self: ask Qiyu about this line (previously was using packed indices, I think it should be unpacked indices).

[deepakn94]

/claude review

[claude]

LGTM — the refactoring is consistent: param_index_map now always stores unpacked offsets, packed offsets are correctly isolated in nvfp4_packed_param_index_map, and all callers (including distrib_optimizer.py) are updated. The removal of get_unpacked_index_map() is clean.

Minor note: tests/unit_tests/distributed/test_param_and_grad_buffer.py has no NVFP4-specific test cases, so this refactoring is not directly exercised by unit tests. Consider adding a test that validates the index maps for a buffer with NVFP4 params (e.g., verifying that param_index_map stores unpacked offsets and nvfp4_packed_param_index_map stores packed offsets).

[deepakn94]

/claude review

[claude]

Good refactoring — the semantics of param_index_map are much cleaner now, and the test coverage is solid.

One test coverage gap: all NVFP4 test cases either mark all params as NVFP4 or none. In practice, NVFP4 buffers often contain a mix (e.g., NVFP4 linear weights alongside bf16 layernorm/embedding weights). A mixed test would exercise the code path where packed and unpacked offsets diverge differently for different params in the same buffer.

The _make_buffer helper already supports this via nvfp4_param_indices. A test like this would fill the gap:

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

[deepakn94]

/claude review

[claude]

LGTM

[deepakn94]

/claude review

[deepakn94]

/claude review

[claude]

LGTM

[svcnvidia-nemo-ci]

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/24613605144