Skip to content

Fix shape metadata when copying tensors across devices - #4245

Open
MrCapricornLiu wants to merge 1 commit into
huggingface:mainfrom
MrCapricornLiu:MrCapricornLiu/fix/inference-tensor-shape-metadata
Open

Fix shape metadata when copying tensors across devices#4245
MrCapricornLiu wants to merge 1 commit into
huggingface:mainfrom
MrCapricornLiu:MrCapricornLiu/fix/inference-tensor-shape-metadata

Conversation

@MrCapricornLiu

Copy link
Copy Markdown

copy_tensor_to_devices, used by pipeline inference with gather_output=True, reduces a shape buffer whose unused entries and non-source ranks are uninitialized. Those values can become dimensions or an invalid dtype. Filtering metadata with nonzero() also loses zero-sized axes, and the receiver passes a tensor where torch.zeros expects a size tuple.

Initialize the metadata buffer to zero, prefix the dimensions with their count and dtype, and construct receiver tensors from the decoded size tuple. The existing collective and dtype mapping are retained.

Tests cover deterministic uninitialized-memory filling, scalars, empty dimensions, three dtypes, and either the first or last rank as source. All 12 metadata subcases fail on the original code and pass with the fix; the utilities module reports 46 passed and 5 skipped. The distributed operations script passes on eight H800 GPUs, and an eight-stage prepare_pippy(gather_output=True) model matches the unpartitioned reference on every rank. Ruff lint/format and pre-commit pass. XLA and multi-node inference were not tested.

Signed-off-by: Chenghao Liu <chliu@stu.pku.edu.cn>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant