From 5589ad3ef64e8844ac1d9b430b2e4de40f596d8a Mon Sep 17 00:00:00 2001 From: RunguoLi Date: Mon, 21 Sep 2026 19:34:18 -0500 Subject: [PATCH] [tensor] keep detach/clone of a padded tensor unpaddable Fixes #6253. `detach()` and `clone()` of a padded tensor copy the padding metadata to the new tensor but do not hijack its own `detach`/`clone`, so the copy looks like a padded tensor to `is_padded_tensor` while `to_unpadded_tensor` fails on it with AttributeError: 'Tensor' object has no attribute '_unpad_detach' `HybridParallelCheckpointIO` always saves `param.detach()`, so `booster.save_model` crashes with `HybridParallelPlugin(tp_size=1)` (any `pp_size`, sharded or not) for every model whose vocab size is not a multiple of `make_vocab_size_divisible_by`, e.g. GPT-2 (50257) or OPT, whose embedding and lm_head are wrapped in `PaddingEmbedding` / `PaddingLMHead`. With `tp_size > 1` the parameter is a distributed tensor and takes a different path, which is why the existing checkpoint tests (LLaMA, vocab 32000) never hit it. Make the copies full padded tensors. Add a unit test for detach/clone + unpad, and a HybridParallelPlugin save/load round trip with a padded vocab (tp=1 with pp=1 and pp=2, sharded and unsharded) that also checks the checkpoint holds the unpadded vocab size. --- colossalai/tensor/padded_tensor/api.py | 3 ++ ...st_hybrid_parallel_plugin_checkpoint_io.py | 33 +++++++++++++++++++ tests/test_tensor/test_padded_tensor.py | 15 +++++++++ 3 files changed, 51 insertions(+) diff --git a/colossalai/tensor/padded_tensor/api.py b/colossalai/tensor/padded_tensor/api.py index 5b66c016b399..d694ab08fc88 100644 --- a/colossalai/tensor/padded_tensor/api.py +++ b/colossalai/tensor/padded_tensor/api.py @@ -14,11 +14,13 @@ def _hijack_detach_and_clone(ptensor: torch.Tensor) -> torch.Tensor: ptensor._unpad_detach = ptensor.detach ptensor._unpad_clone = ptensor.clone + # the copies must be full padded tensors too, otherwise to_unpadded_tensor() fails on them def new_detach(self): t_ = self._unpad_detach() t_._padding_dim = self._padding_dim t_._origin_length = self._origin_length t_._current_length = self._current_length + _hijack_detach_and_clone(t_) return t_ def new_clone(self, *args, **kwargs): @@ -26,6 +28,7 @@ def new_clone(self, *args, **kwargs): t_._padding_dim = self._padding_dim t_._origin_length = self._origin_length t_._current_length = self._current_length + _hijack_detach_and_clone(t_) return t_ # bind the new methods to the tensor diff --git a/tests/test_checkpoint_io/test_hybrid_parallel_plugin_checkpoint_io.py b/tests/test_checkpoint_io/test_hybrid_parallel_plugin_checkpoint_io.py index a338d98f4746..eb728779e543 100644 --- a/tests/test_checkpoint_io/test_hybrid_parallel_plugin_checkpoint_io.py +++ b/tests/test_checkpoint_io/test_hybrid_parallel_plugin_checkpoint_io.py @@ -1,8 +1,11 @@ +import glob + import pytest import torch import torch.distributed as dist from packaging.version import Version from torch.optim import Adam +from transformers import GPT2Config, GPT2LMHeadModel from utils import shared_tempdir import colossalai @@ -143,9 +146,39 @@ def _preprocess_data(data): clear_layout_converter() +@clear_cache_before_run() +@parameterize("shard", [True, False]) +@parameterize("test_config", [{"tp_size": 1, "pp_size": 1}, {"tp_size": 1, "pp_size": 2, "num_microbatches": 2}]) +def exam_state_dict_padded_vocab(shard: bool, test_config: dict): + # a vocab size that is not a multiple of make_vocab_size_divisible_by (64), so the embedding and + # lm_head weights are padded tensors, which must be unpadded when saving (#6253) + config = GPT2Config(n_layer=2, n_head=4, n_embd=64, vocab_size=1000, n_positions=64) + booster = Booster(plugin=HybridParallelPlugin(**test_config, precision="fp16", initial_scale=1)) + model = GPT2LMHeadModel(config).cuda() + model, *_ = booster.boost(model, Adam(model.parameters(), lr=1e-3)) + + with shared_tempdir() as tempdir: + model_ckpt_path = f"{tempdir}/model" + booster.save_model(model, model_ckpt_path, shard=shard) + dist.barrier() + if dist.get_rank() == 0: + # the checkpoint must hold the original (unpadded) vocab size + files = glob.glob(f"{model_ckpt_path}/*.bin") if shard else [model_ckpt_path] + saved = {} + for f in files: + saved.update(torch.load(f, map_location="cpu")) + assert saved["transformer.wte.weight"].shape == (config.vocab_size, config.n_embd) + new_model = GPT2LMHeadModel(config).cuda() + new_model, *_ = booster.boost(new_model, Adam(new_model.parameters(), lr=1e-3)) + booster.load_model(new_model, model_ckpt_path) + check_state_dict_equal(model.unwrap().state_dict(), new_model.unwrap().state_dict()) + dist.barrier() + + def run_dist(rank, world_size, port): colossalai.launch(rank=rank, world_size=world_size, host="localhost", port=port, backend="nccl") exam_state_dict() + exam_state_dict_padded_vocab() @pytest.mark.dist diff --git a/tests/test_tensor/test_padded_tensor.py b/tests/test_tensor/test_padded_tensor.py index 6d19845dff2f..a6062b93d281 100644 --- a/tests/test_tensor/test_padded_tensor.py +++ b/tests/test_tensor/test_padded_tensor.py @@ -36,6 +36,21 @@ def check_padded_tensor(rank, world_size, port): assert global_tensor.shape == original_tensor.shape +def test_unpad_detached_and_cloned_tensor(): + # detach() / clone() of a padded tensor must give a padded tensor that can be unpadded on its own, + # this is what checkpoint saving does before writing a parameter + original = torch.rand(10, 4) + padded = to_padded_tensor(original.clone(), current_length=16, padding_dim=0) + for tensor_copy in (padded.detach(), padded.clone(), padded.detach().clone()): + assert is_padded_tensor(tensor_copy) + assert tensor_copy.shape == (16, 4) + unpadded = to_unpadded_tensor(tensor_copy) + assert not is_padded_tensor(unpadded) + assert torch.equal(unpadded, original) + # the source tensor is still padded + assert is_padded_tensor(padded) and padded.shape == (16, 4) + + @rerun_if_address_is_in_use() def test_padded_tensor(): world_size = 4