Add _TensorViewCodec for storage-deduplicated tensor serialization

[fzyzcjy]

Introduce _TensorViewCodec (in checkpoint_transfer.py), which encodes a list
of tensors as (unique_storages, view_metas) by deduping shared underlying
storages (e.g. Megatron distributed-optimizer grad buckets) and reconstructs the
original views via as_strided. Comprehensively unit-tested by
TestTensorViewCodec. Consumed by the peer checkpoint transfer added next.

[gemini-code-assist]

Code Review

This pull request introduces _TensorViewCodec to encode and decode PyTorch tensors sharing the same underlying storage, optimizing memory usage during checkpoint transfers. It also includes comprehensive unit tests to verify correctness across various scenarios. The feedback suggests two important improvements: first, to include the device in the deduplication key to prevent storage ID collisions across different devices; second, to slice the untyped storage to a multiple of the target element size during decoding to avoid potential RuntimeErrors when reinterpreting the dtype.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

Using only storage.data_ptr() as the key in storage_id_by_ptr can lead to collisions if the input list contains tensors from different devices (e.g., CPU and GPU, or multiple GPUs). Since raw memory addresses are independent across devices, they can easily overlap. To prevent incorrect storage deduplication, use a tuple of (t.device, ptr) as the dictionary key.

Suggested change

	storage_id_by_ptr: dict[int, int] = {}
	unique_storages: list[torch.Tensor] = []
	view_metas: list[dict] = []
	for t in tensors:
	storage = t.untyped_storage()
	ptr = storage.data_ptr()
	if ptr not in storage_id_by_ptr:
	storage_id_by_ptr[ptr] = len(unique_storages)
	# Wrap full storage as uint8 tensor (no copy, shares memory).
	unique_storages.append(torch.tensor(storage, dtype=torch.uint8, device=t.device))
	view_metas.append(
	{
	"storage_id": storage_id_by_ptr[ptr],
	"dtype": t.dtype,
	"shape": tuple(t.shape),
	"stride": tuple(t.stride()),
	"storage_offset": t.storage_offset(),
	}
	)
	storage_id_by_ptr: dict[tuple[torch.device, int], int] = {}
	unique_storages: list[torch.Tensor] = []
	view_metas: list[dict] = []
	for t in tensors:
	storage = t.untyped_storage()
	ptr = storage.data_ptr()
	key = (t.device, ptr)
	if key not in storage_id_by_ptr:
	storage_id_by_ptr[key] = len(unique_storages)
	# Wrap full storage as uint8 tensor (no copy, shares memory).
	unique_storages.append(torch.tensor(storage, dtype=torch.uint8, device=t.device))
	view_metas.append(
	{
	"storage_id": storage_id_by_ptr[key],
	"dtype": t.dtype,
	"shape": tuple(t.shape),
	"stride": tuple(t.stride()),
	"storage_offset": t.storage_offset(),
	}
	)

[Shi-Dong]

I think this concern is legitimate and worth addressing.

[gemini-code-assist]

If the untyped storage size is not a multiple of the element size of vm["dtype"] (which can happen with mixed-dtype storages or padded storages), calling storage_t.view(vm["dtype"]) will raise a RuntimeError. To prevent this, slice storage_t to a multiple of the element size before reinterpreting its dtype.

Suggested change

	for vm in view_metas:
	storage_t = unique_storages[vm["storage_id"]] # uint8 view of received storage
	# Reinterpret bytes as the original dtype, then apply stride/offset.
	dtype_view = storage_t.view(vm["dtype"])
	view = torch.as_strided(
	dtype_view,
	size=vm["shape"],
	stride=vm["stride"],
	storage_offset=vm["storage_offset"],
	)
	tensors.append(view)
	for vm in view_metas:
	storage_t = unique_storages[vm["storage_id"]] # uint8 view of received storage
	# Reinterpret bytes as the original dtype, then apply stride/offset.
	element_size = torch.tensor([], dtype=vm["dtype"]).element_size()
	num_bytes = (storage_t.numel() // element_size) * element_size
	dtype_view = storage_t[:num_bytes].view(vm["dtype"])
	view = torch.as_strided(
	dtype_view,
	size=vm["shape"],
	stride=vm["stride"],
	storage_offset=vm["storage_offset"],
	)
	tensors.append(view)

[Shi-Dong]

Please address the Gemini concern on duplicated dict keys.

[Shi-Dong]

I think this concern is legitimate and worth addressing.