From b7a88ce3bebe3c3b24d6c0d34ed65dc9a6001a4a Mon Sep 17 00:00:00 2001 From: wangzhijian Date: Thu, 22 Oct 2026 07:07:21 +0000 Subject: [PATCH 1/2] [ci]: upgrade Transformers to 5.17 Align PEFT, Hub, tokenizers and Diffusers dependencies with Transformers 5.17. Update the CI image profile and add offline CPU LoRA and diffusion pipeline import checks. --- docker/torch213/Dockerfile | 7 +-- docker/torch213/README.md | 57 +++++++++++++++++++++-- docker/torch213/requirements.lock | 14 +++--- docker/torch213/verify_image.py | 77 ++++++++++++++++++++++++++++++- requirements/requirements.txt | 8 ++-- 5 files changed, 144 insertions(+), 19 deletions(-) diff --git a/docker/torch213/Dockerfile b/docker/torch213/Dockerfile index c977b552e2da..056df69a7544 100644 --- a/docker/torch213/Dockerfile +++ b/docker/torch213/Dockerfile @@ -43,12 +43,13 @@ RUN --mount=type=bind,source=docker/torch213/wheelhouse,target=/wheelhouse,reado # resurrect their old .dist-info entries alongside the locked versions. FROM dependencies AS ci ARG SOURCE_COMMIT=unknown -LABEL org.opencontainers.image.title="ColossalAI Torch 2.13 CI dependencies" \ +LABEL org.opencontainers.image.title="ColossalAI Torch 2.13 / Transformers 5.17 CI dependencies" \ org.opencontainers.image.source="https://github.com/hpcaitech/ColossalAI" \ org.opencontainers.image.revision="${SOURCE_COMMIT}" COPY docker/torch213/requirements*.lock docker/torch213/verify_image.py \ docker/torch213/Dockerfile docker/torch213/Dockerfile.dockerignore /opt/ci/ WORKDIR /workspace -# This checks metadata, imports and CPU operations; CUDA remains hidden. -RUN NVIDIA_VISIBLE_DEVICES=void CUDA_VISIBLE_DEVICES="" python /opt/ci/verify_image.py +# Check metadata, CPU operations and tiny LoRA training without GPU/model downloads. +RUN NVIDIA_VISIBLE_DEVICES=void CUDA_VISIBLE_DEVICES="" \ + HF_HUB_OFFLINE=1 HF_DATASETS_OFFLINE=1 python /opt/ci/verify_image.py CMD ["/bin/bash"] diff --git a/docker/torch213/README.md b/docker/torch213/README.md index e5b607315517..608f58a36d52 100644 --- a/docker/torch213/README.md +++ b/docker/torch213/README.md @@ -1,4 +1,4 @@ -# Torch 2.13 / CUDA 13.0 CI dependency image +# Torch 2.13 / CUDA 13.0 / Transformers 5.17 CI dependency image This image provides an isolated Python environment and CUDA build tools for the existing layered GPU tests. It contains **no ColossalAI checkout, installed @@ -21,7 +21,9 @@ and actual container/GPU qualification are separate acceptance steps. | torchvision | 0.28.0+cu130 | | torchaudio | 2.11.0+cu130 (existing development artifact; see below) | | Triton | 3.7.1 | -| Transformers / PEFT | 4.51.3 / 0.13.2 | +| Transformers / PEFT | 5.17.0 / 0.19.1 | +| Hugging Face Hub / Tokenizers | 1.5.0 / 0.23.1 | +| Diffusers | 0.35.2 | | pytest / pytest-testmon | 7.4.4 / 2.2.0 | | setuptools | 84.0.0 | | PyTorch's Python cuDNN package | nvidia-cudnn-cu13==9.20.0.48 | @@ -53,12 +55,57 @@ Report skipped coverage separately; this image is not a claim of passing every optional test. The general `requirements/requirements-test.txt` still includes flash-attn, so do not reinstall it blindly on top of this locked core profile. -The library permits `torch>=2.2.0,<2.14`, retaining the 2.5.1 CPU baseline. That +The library permits `torch>=2.5,<2.14`, matching Transformers 5.17's minimum Torch +requirement and retaining the 2.5.1 CPU baseline. That installation range is distinct from this exact GPU profile and does not imply -that every intermediate Torch/CUDA combination has been tested. Transformers -remains at 4.51.3. On Torch 2.13 the FP8 implementation currently uses its eager +that every intermediate Torch/CUDA combination has been tested. Transformers and +PEFT are pinned to 5.17.0 and 0.19.1 respectively. The image verifier runs two +offline CPU optimizer steps on a tiny Llama with LoRA, checking finite losses, +finite nonzero gradients and actual adapter updates. This is a dependency smoke +check, not a substitute for ColossalAI's GPU/distributed regressions. On Torch +2.13 the FP8 implementation currently uses its eager fallback; passing FP8 tests does not establish `torch.compile` support. +Diffusers is upgraded from 0.29.0 to 0.35.2 because the older version imports +`FLAX_WEIGHTS_NAME` from Transformers when resolving `DiffusionPipeline`; that +symbol is no longer available in Transformers 5.17. The image verifier also +resolves the DiffusionPipeline, PixArt Alpha and Stable Diffusion 3 lazy imports, +which `pip check` and a plain `import diffusers` do not exercise. These checks +do not establish numerical correctness of diffusion inference. + +Do not switch the shared `COLOSSAL_CI_IMAGE` variable to a `docker commit` +experiment that still embeds the previous `/opt/ci` recipe. Build and qualify +a fresh candidate from these inputs first; source requirements, image locks and +embedded recipe must agree. Keep the stable image available for rollback, and +coordinate changing the shared image variable with merging the dependency PR. + +## Small offline source checks + +Inside a development container using this dependency profile, install the source +without replacing the locked packages, then run a CPU-only HF5 subset: + +```bash +# Install the checkout, without compiling all optional CUDA extensions up front. +BUILD_EXT=0 python -m pip install --no-deps --no-build-isolation -e . +python -m pip check + +# These selected tests do not use FAST_TEST; do not confuse them with full GPU CI. +unset FAST_TEST +HF_HUB_OFFLINE=1 HF_DATASETS_OFFLINE=1 OMP_NUM_THREADS=2 MKL_NUM_THREADS=2 \ + python -m pytest -q -ra \ + tests/test_shardformer/test_hf5_compat/test_cpu_rng.py \ + tests/test_shardformer/test_hf5_compat/test_integrated_core.py \ + tests/test_shardformer/test_hf5_compat/test_integration_alias.py \ + tests/test_shardformer/test_hf5_compat/test_integration_loading.py \ + tests/test_shardformer/test_hf5_compat/test_qwen_loss.py +``` + +This covers HF5 policy dispatch, small non-TP/EP training, checkpoint/loading +contracts and single-rank CPU Gloo Qwen losses. It does not validate CUDA kernels, +tensor/expert parallelism, multi-GPU communication or the Torch 2.5.1 CPU baseline. +The complete image verifier must run **before** installing ColossalAI: it +deliberately rejects images containing an installed ColossalAI distribution. + ## Build Use a dedicated or explicitly approved build host. A normal Docker build/load/run diff --git a/docker/torch213/requirements.lock b/docker/torch213/requirements.lock index e235e03a08ef..dc658e19c425 100644 --- a/docker/torch213/requirements.lock +++ b/docker/torch213/requirements.lock @@ -1,5 +1,5 @@ # Dependency snapshot for Linux x86_64, CPython 3.10 and requirements-torch.lock. -# Captured from the existing Torch 2.13 development environment. No ColossalAI, +# Torch 2.13 / Transformers 5.17 dependency profile. No ColossalAI, # personal paths, VCS dependencies or optional GPU backends are included. # All entries are installed with --no-deps and checked with pip check. pip==26.2.1 @@ -31,7 +31,7 @@ cuda-toolkit==13.0.3.0 datasets==5.0.1 decorator==5.3.1 Deprecated==1.3.1 -diffusers==0.29.0 +diffusers==0.35.2 dill==0.4.1 distlib==0.4.3 einops==0.8.2 @@ -47,7 +47,7 @@ h11==0.16.0 hf-xet==1.6.0 httpcore==1.0.9 httpx==0.28.1 -huggingface_hub==0.36.2 +huggingface_hub==1.5.0 identify==2.6.19 idna==3.19 importlib_metadata==9.0.1 @@ -85,7 +85,7 @@ nvidia-nvtx==13.0.85 packaging==26.3 pandas==2.3.3 paramiko==5.0.0 -peft==0.13.2 +peft==0.19.1 pillow==12.3.0 platformdirs==4.11.8 pluggy==1.6.0 @@ -115,15 +115,17 @@ rpds-py==0.30.0 rpyc==6.0.0 safetensors==0.8.0 sentencepiece==0.2.2 +shellingham==1.5.4 six==1.17.0 soupsieve==2.9.2 starlette==1.6.0 sympy==1.14.0 timm==1.0.29 -tokenizers==0.21.4 +tokenizers==0.23.1 tomli==2.4.1 tqdm==4.70.1 -transformers==4.51.3 +transformers==5.17.0 +typer==0.27.3 typing-inspection==0.4.4 typing_extensions==4.16.0 tzdata==2026.4 diff --git a/docker/torch213/verify_image.py b/docker/torch213/verify_image.py index 01ed90e15fd8..c95e4763b2c2 100644 --- a/docker/torch213/verify_image.py +++ b/docker/torch213/verify_image.py @@ -12,6 +12,68 @@ from packaging.version import Version +def verify_transformers_peft(): + """Exercise offline CPU LoRA training, not just the packages' lazy imports.""" + import torch + from peft import LoraConfig, get_peft_model + from transformers import LlamaConfig, LlamaForCausalLM + + previous_threads = torch.get_num_threads() + try: + torch.set_num_threads(2) + with torch.random.fork_rng(devices=[]): + torch.manual_seed(42) + config = LlamaConfig( + vocab_size=64, + hidden_size=32, + intermediate_size=64, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + max_position_embeddings=64, + attn_implementation="eager", + use_cache=False, + ) + model = get_peft_model( + LlamaForCausalLM(config).to(device="cpu", dtype=torch.float32), + LoraConfig(task_type="CAUSAL_LM", r=4, lora_alpha=8, target_modules=["q_proj", "v_proj"]), + ) + model.train() + trainable = [parameter for parameter in model.parameters() if parameter.requires_grad] + assert trainable, "LoRA has no trainable parameters" + assert all(parameter.device.type == "cpu" for parameter in model.parameters()) + before = [parameter.detach().clone() for parameter in trainable] + optimizer = torch.optim.AdamW(trainable, lr=1e-3) + tokens = torch.randint(0, config.vocab_size, (2, 16), device="cpu") + losses = [] + for _ in range(2): + optimizer.zero_grad(set_to_none=True) + loss = model(input_ids=tokens, labels=tokens).loss + assert torch.isfinite(loss), "LoRA loss is not finite" + loss.backward() + gradients = [parameter.grad for parameter in trainable] + assert all(gradient is not None and torch.isfinite(gradient).all() for gradient in gradients) + assert any( + torch.count_nonzero(gradient).item() for gradient in gradients + ), "LoRA gradients are all zero" + optimizer.step() + losses.append(loss.item()) + assert any(not torch.equal(old, new) for old, new in zip(before, trainable)), "LoRA weights did not update" + return {"device": "cpu", "steps": len(losses), "losses": losses} + finally: + torch.set_num_threads(previous_threads) + + +def verify_diffusers_pipelines(): + """Resolve lazy pipeline imports used by ColossalAI, without downloading models.""" + import diffusers + + pipelines = ("DiffusionPipeline", "PixArtAlphaPipeline", "StableDiffusion3Pipeline") + for name in pipelines: + assert issubclass(getattr(diffusers, name), diffusers.DiffusionPipeline) + return list(pipelines) + + def main(): assert sys.version_info[:2] == (3, 10), sys.version root = Path(__file__).resolve().parent @@ -53,9 +115,22 @@ def main(): assert torchvision.ops.nms(boxes, torch.tensor([0.9, 0.8]), 0.5).tolist() == [0] waveform = torch.zeros(1, 160) assert torchaudio.functional.resample(waveform, 16000, 8000).shape == (1, 80) + hf_smoke = verify_transformers_peft() + diffusion_imports = verify_diffusers_pipelines() # pytest --help loads installed entry points, including pytest-testmon. subprocess.run([sys.executable, "-m", "pytest", "--help"], check=True, stdout=subprocess.DEVNULL) - print(json.dumps({"python": sys.version, "torch_cuda": torch.version.cuda, "packages": versions}, indent=2)) + print( + json.dumps( + { + "python": sys.version, + "torch_cuda": torch.version.cuda, + "packages": versions, + "hf_smoke": hf_smoke, + "diffusion_pipeline_imports": diffusion_imports, + }, + indent=2, + ) + ) if __name__ == "__main__": diff --git a/requirements/requirements.txt b/requirements/requirements.txt index c84a28969484..dffadb0afc95 100644 --- a/requirements/requirements.txt +++ b/requirements/requirements.txt @@ -7,7 +7,7 @@ click fabric contexttimer ninja -torch>=2.2.0,<2.14 +torch>=2.5,<2.14 safetensors einops pydantic @@ -15,11 +15,11 @@ ray sentencepiece google protobuf -transformers==4.51.3 -peft>=0.7.1,<=0.13.2 +transformers==5.17.0 +peft==0.19.1 bitsandbytes>=0.39.0 rpyc==6.0.0 fastapi uvicorn==0.29.0 galore_torch -diffusers==0.29.0 +diffusers==0.35.2 From cc471e82f68c142594e0dffd4b9347613129c929 Mon Sep 17 00:00:00 2001 From: wangzhijian Date: Thu, 22 Oct 2026 17:19:58 +0800 Subject: [PATCH 2/2] [compatibility]: fix ChatGLM and PEFT for Transformers 5.17 Preserve ChatGLM configuration fields removed from Transformers defaults. Upgrade PEFT to 0.21.2 and align the CI image dependency profile. Add offline ChatGLM and distributed LoRA checkpoint regressions to CPU CI and image validation. --- .github/scripts/ci/run_colossalai_batch.sh | 2 + .github/scripts/ci/run_cpu_unit_tests.sh | 2 + .../chatglm2_6b/configuration_chatglm.py | 4 + .../modeling/chatglm2_6b/modeling_chatglm.py | 4 +- docker/torch213/Dockerfile | 2 +- docker/torch213/README.md | 25 +++++-- docker/torch213/requirements.lock | 2 +- docker/torch213/test_verify_image.py | 16 ++++ docker/torch213/verify_image.py | 74 ++++++++++++++++++- requirements/requirements.txt | 2 +- tests/test_lazy/test_chatglm_cpu.py | 73 ++++++++++++++++++ 11 files changed, 193 insertions(+), 13 deletions(-) create mode 100644 docker/torch213/test_verify_image.py create mode 100644 tests/test_lazy/test_chatglm_cpu.py diff --git a/.github/scripts/ci/run_colossalai_batch.sh b/.github/scripts/ci/run_colossalai_batch.sh index a419e5d91d1e..01db482cca27 100644 --- a/.github/scripts/ci/run_colossalai_batch.sh +++ b/.github/scripts/ci/run_colossalai_batch.sh @@ -71,8 +71,10 @@ case "${BATCH}" in EXPECTED_GPUS=0 DEFAULT_TIMEOUT_MIN=15 TESTS=( + docker/torch213/test_verify_image.py tests/test_config/test_load_config.py tests/test_infer/test_async_engine/test_request_tracer.py + tests/test_lazy/test_chatglm_cpu.py tests/test_optimizer/test_lr_scheduler.py tests/test_pipeline/test_pipeline_utils/test_t5_pipeline_utils.py tests/test_pipeline/test_pipeline_utils/test_whisper_pipeline_utils.py diff --git a/.github/scripts/ci/run_cpu_unit_tests.sh b/.github/scripts/ci/run_cpu_unit_tests.sh index 8d39713f95ca..2677eccf2d84 100644 --- a/.github/scripts/ci/run_cpu_unit_tests.sh +++ b/.github/scripts/ci/run_cpu_unit_tests.sh @@ -9,8 +9,10 @@ PYTHON="${COLOSSAL_PYTHON:-python}" TESTS=( .github/workflows/scripts/example_checks/test_run_test_ci.py + docker/torch213/test_verify_image.py tests/test_config/test_load_config.py tests/test_infer/test_async_engine/test_request_tracer.py + tests/test_lazy/test_chatglm_cpu.py tests/test_optimizer/test_lr_scheduler.py tests/test_pipeline/test_pipeline_utils/test_t5_pipeline_utils.py tests/test_pipeline/test_pipeline_utils/test_whisper_pipeline_utils.py diff --git a/colossalai/shardformer/modeling/chatglm2_6b/configuration_chatglm.py b/colossalai/shardformer/modeling/chatglm2_6b/configuration_chatglm.py index bb774676a4d4..aa0b614bf5e7 100644 --- a/colossalai/shardformer/modeling/chatglm2_6b/configuration_chatglm.py +++ b/colossalai/shardformer/modeling/chatglm2_6b/configuration_chatglm.py @@ -30,6 +30,7 @@ def __init__( quantization_bit=0, pre_seq_len=None, prefix_projection=False, + use_cache=True, **kwargs, ): self.num_layers = num_layers @@ -58,3 +59,6 @@ def __init__( self.pre_seq_len = pre_seq_len self.prefix_projection = prefix_projection super().__init__(**kwargs) + # Transformers 5 no longer preserves this generation option from kwargs. + # It is also part of ChatGLM's forward contract, so own it explicitly. + self.use_cache = use_cache diff --git a/colossalai/shardformer/modeling/chatglm2_6b/modeling_chatglm.py b/colossalai/shardformer/modeling/chatglm2_6b/modeling_chatglm.py index 6ae4b06e517a..b64544f50c75 100644 --- a/colossalai/shardformer/modeling/chatglm2_6b/modeling_chatglm.py +++ b/colossalai/shardformer/modeling/chatglm2_6b/modeling_chatglm.py @@ -1004,7 +1004,9 @@ class ChatGLMForConditionalGeneration(ChatGLMPreTrainedModel): def __init__(self, config: ChatGLMConfig, empty_init=True, device=None): super().__init__(config) - self.max_sequence_length = config.max_length + # Transformers 5 no longer supplies generation defaults on PretrainedConfig. + # Preserve an explicit legacy value, otherwise use ChatGLM's context length. + self.max_sequence_length = getattr(config, "max_length", config.seq_length) self.transformer = ChatGLMModel(config, empty_init=empty_init, device=device) self.config = config self.quantized = False diff --git a/docker/torch213/Dockerfile b/docker/torch213/Dockerfile index 056df69a7544..4f4b3909e8e6 100644 --- a/docker/torch213/Dockerfile +++ b/docker/torch213/Dockerfile @@ -49,7 +49,7 @@ LABEL org.opencontainers.image.title="ColossalAI Torch 2.13 / Transformers 5.17 COPY docker/torch213/requirements*.lock docker/torch213/verify_image.py \ docker/torch213/Dockerfile docker/torch213/Dockerfile.dockerignore /opt/ci/ WORKDIR /workspace -# Check metadata, CPU operations and tiny LoRA training without GPU/model downloads. +# Check metadata, CPU operations and LoRA train/save/reload without GPU/model downloads. RUN NVIDIA_VISIBLE_DEVICES=void CUDA_VISIBLE_DEVICES="" \ HF_HUB_OFFLINE=1 HF_DATASETS_OFFLINE=1 python /opt/ci/verify_image.py CMD ["/bin/bash"] diff --git a/docker/torch213/README.md b/docker/torch213/README.md index 608f58a36d52..af44ef49ffe6 100644 --- a/docker/torch213/README.md +++ b/docker/torch213/README.md @@ -21,7 +21,7 @@ and actual container/GPU qualification are separate acceptance steps. | torchvision | 0.28.0+cu130 | | torchaudio | 2.11.0+cu130 (existing development artifact; see below) | | Triton | 3.7.1 | -| Transformers / PEFT | 5.17.0 / 0.19.1 | +| Transformers / PEFT | 5.17.0 / 0.21.2 | | Hugging Face Hub / Tokenizers | 1.5.0 / 0.23.1 | | Diffusers | 0.35.2 | | pytest / pytest-testmon | 7.4.4 / 2.2.0 | @@ -59,10 +59,19 @@ The library permits `torch>=2.5,<2.14`, matching Transformers 5.17's minimum Tor requirement and retaining the 2.5.1 CPU baseline. That installation range is distinct from this exact GPU profile and does not imply that every intermediate Torch/CUDA combination has been tested. Transformers and -PEFT are pinned to 5.17.0 and 0.19.1 respectively. The image verifier runs two +PEFT are pinned to 5.17.0 and 0.21.2 respectively. The image verifier runs two offline CPU optimizer steps on a tiny Llama with LoRA, checking finite losses, -finite nonzero gradients and actual adapter updates. This is a dependency smoke -check, not a substitute for ColossalAI's GPU/distributed regressions. On Torch +finite nonzero gradients and actual adapter updates. It then saves and reloads +both `.bin` and safetensors adapters, checks adapter weights and evaluation +logits, and verifies another optimizer update after loading. These checks run +both without a process group and with a single-rank CPU Gloo group. PEFT 0.19.1 +attempts to import the removed `EmbeddingParallel` symbol only when a process +group is initialized, so ordinary single-process training and reload checks +alone miss this regression. PEFT 0.21.2 includes the upstream TP compatibility +fixes; the image must be rebuilt when updating this pin. + +This is a dependency smoke check, not a substitute for ColossalAI's +GPU/distributed regressions. On Torch 2.13 the FP8 implementation currently uses its eager fallback; passing FP8 tests does not establish `torch.compile` support. @@ -93,6 +102,8 @@ python -m pip check unset FAST_TEST HF_HUB_OFFLINE=1 HF_DATASETS_OFFLINE=1 OMP_NUM_THREADS=2 MKL_NUM_THREADS=2 \ python -m pytest -q -ra \ + docker/torch213/test_verify_image.py \ + tests/test_lazy/test_chatglm_cpu.py \ tests/test_shardformer/test_hf5_compat/test_cpu_rng.py \ tests/test_shardformer/test_hf5_compat/test_integrated_core.py \ tests/test_shardformer/test_hf5_compat/test_integration_alias.py \ @@ -100,8 +111,10 @@ HF_HUB_OFFLINE=1 HF_DATASETS_OFFLINE=1 OMP_NUM_THREADS=2 MKL_NUM_THREADS=2 \ tests/test_shardformer/test_hf5_compat/test_qwen_loss.py ``` -This covers HF5 policy dispatch, small non-TP/EP training, checkpoint/loading -contracts and single-rank CPU Gloo Qwen losses. It does not validate CUDA kernels, +This covers ChatGLM initialization/lazy materialization, LoRA checkpoint +roundtrips, HF5 policy dispatch, small non-TP/EP training, checkpoint/loading +contracts and single-rank CPU Gloo Qwen losses. The new ChatGLM and LoRA checks +also run in the Layer 1 and fork-safe CPU gates. This does not validate CUDA kernels, tensor/expert parallelism, multi-GPU communication or the Torch 2.5.1 CPU baseline. The complete image verifier must run **before** installing ColossalAI: it deliberately rejects images containing an installed ColossalAI distribution. diff --git a/docker/torch213/requirements.lock b/docker/torch213/requirements.lock index dc658e19c425..983f54fbf56f 100644 --- a/docker/torch213/requirements.lock +++ b/docker/torch213/requirements.lock @@ -85,7 +85,7 @@ nvidia-nvtx==13.0.85 packaging==26.3 pandas==2.3.3 paramiko==5.0.0 -peft==0.19.1 +peft==0.21.2 pillow==12.3.0 platformdirs==4.11.8 pluggy==1.6.0 diff --git a/docker/torch213/test_verify_image.py b/docker/torch213/test_verify_image.py new file mode 100644 index 000000000000..d9aa76470482 --- /dev/null +++ b/docker/torch213/test_verify_image.py @@ -0,0 +1,16 @@ +import torch +import torch.distributed as dist +from verify_image import verify_transformers_peft + + +def test_cpu_lora_checkpoint_roundtrip(): + """The image smoke check must exercise local and distributed adapter reloads.""" + threads = torch.get_num_threads() + rng_state = torch.random.get_rng_state().clone() + result = verify_transformers_peft() + assert result["device"] == "cpu" + assert result["steps"] == 2 + assert result["checkpoint_checks"] == ["local-bin", "local-safetensors", "gloo-bin", "gloo-safetensors"] + assert not dist.is_initialized() + assert torch.get_num_threads() == threads + assert torch.equal(torch.random.get_rng_state(), rng_state) diff --git a/docker/torch213/verify_image.py b/docker/torch213/verify_image.py index c95e4763b2c2..fe3a06f574f7 100644 --- a/docker/torch213/verify_image.py +++ b/docker/torch213/verify_image.py @@ -6,14 +6,80 @@ import shutil import subprocess import sys +from copy import deepcopy +from datetime import timedelta from pathlib import Path +from tempfile import TemporaryDirectory from packaging.requirements import Requirement from packaging.version import Version +def verify_lora_checkpoints(model, base_model, tokens): + """Check adapter reload and resumed updates, including the distributed PEFT path.""" + import torch + import torch.distributed as dist + from peft import PeftModel, get_peft_model_state_dict + + assert dist.is_available() and dist.is_gloo_available(), "CPU LoRA checkpoint checks require Gloo" + assert not dist.is_initialized(), "Run image validation outside an existing process group" + model.eval() + with torch.no_grad(): + reference_logits = model(input_ids=tokens).logits + expected_adapter = get_peft_model_state_dict(model) + checked = [] + with TemporaryDirectory(prefix="ci-lora-checkpoint-") as temporary: + for distributed in (False, True): + if distributed: + # PEFT 0.19.1 only takes the broken TP import path when a process + # group exists, even for a non-TP model. One CPU rank reproduces it. + dist.init_process_group( + "gloo", + init_method=Path(temporary, "rendezvous").as_uri(), + rank=0, + world_size=1, + timeout=timedelta(seconds=30), + ) + try: + for safe_serialization in (False, True): + mode = f"{'gloo' if distributed else 'local'}-{'safetensors' if safe_serialization else 'bin'}" + checkpoint = Path(temporary, mode) + model.save_pretrained(checkpoint, safe_serialization=safe_serialization) + expected_file = "adapter_model.safetensors" if safe_serialization else "adapter_model.bin" + assert (checkpoint / expected_file).is_file(), f"Missing {mode} checkpoint" + restored = PeftModel.from_pretrained( + deepcopy(base_model), checkpoint, is_trainable=True, local_files_only=True + ) + actual_adapter = get_peft_model_state_dict(restored) + assert actual_adapter.keys() == expected_adapter.keys(), f"Adapter keys changed for {mode}" + for name, expected in expected_adapter.items(): + torch.testing.assert_close(actual_adapter[name], expected, rtol=0, atol=0) + restored.eval() + with torch.no_grad(): + torch.testing.assert_close(restored(input_ids=tokens).logits, reference_logits) + restored.train() + trainable = [parameter for parameter in restored.parameters() if parameter.requires_grad] + assert trainable, f"Reloaded {mode} adapter is not trainable" + before = [parameter.detach().clone() for parameter in trainable] + optimizer = torch.optim.AdamW(trainable, lr=1e-3) + optimizer.zero_grad(set_to_none=True) + loss = restored(input_ids=tokens, labels=tokens).loss + assert torch.isfinite(loss), f"Non-finite resumed loss for {mode}" + loss.backward() + gradients = [parameter.grad for parameter in trainable] + assert all(gradient is not None and torch.isfinite(gradient).all() for gradient in gradients) + assert any(torch.count_nonzero(gradient).item() for gradient in gradients) + optimizer.step() + assert any(not torch.equal(old, new) for old, new in zip(before, trainable)) + checked.append(mode) + finally: + if distributed: + dist.destroy_process_group() + return checked + + def verify_transformers_peft(): - """Exercise offline CPU LoRA training, not just the packages' lazy imports.""" + """Exercise offline CPU LoRA training, checkpoint reload and resumed training.""" import torch from peft import LoraConfig, get_peft_model from transformers import LlamaConfig, LlamaForCausalLM @@ -34,8 +100,9 @@ def verify_transformers_peft(): attn_implementation="eager", use_cache=False, ) + base_model = LlamaForCausalLM(config).to(device="cpu", dtype=torch.float32) model = get_peft_model( - LlamaForCausalLM(config).to(device="cpu", dtype=torch.float32), + deepcopy(base_model), LoraConfig(task_type="CAUSAL_LM", r=4, lora_alpha=8, target_modules=["q_proj", "v_proj"]), ) model.train() @@ -59,7 +126,8 @@ def verify_transformers_peft(): optimizer.step() losses.append(loss.item()) assert any(not torch.equal(old, new) for old, new in zip(before, trainable)), "LoRA weights did not update" - return {"device": "cpu", "steps": len(losses), "losses": losses} + checkpoints = verify_lora_checkpoints(model, base_model, tokens) + return {"device": "cpu", "steps": len(losses), "losses": losses, "checkpoint_checks": checkpoints} finally: torch.set_num_threads(previous_threads) diff --git a/requirements/requirements.txt b/requirements/requirements.txt index dffadb0afc95..fc6b13e2c1ef 100644 --- a/requirements/requirements.txt +++ b/requirements/requirements.txt @@ -16,7 +16,7 @@ sentencepiece google protobuf transformers==5.17.0 -peft==0.19.1 +peft==0.21.2 bitsandbytes>=0.39.0 rpyc==6.0.0 fastapi diff --git a/tests/test_lazy/test_chatglm_cpu.py b/tests/test_lazy/test_chatglm_cpu.py new file mode 100644 index 000000000000..8879de762f65 --- /dev/null +++ b/tests/test_lazy/test_chatglm_cpu.py @@ -0,0 +1,73 @@ +import pytest +import torch + +from colossalai.lazy import LazyInitContext, LazyTensor +from colossalai.shardformer.modeling.chatglm2_6b.configuration_chatglm import ChatGLMConfig +from colossalai.shardformer.modeling.chatglm2_6b.modeling_chatglm import ChatGLMForConditionalGeneration + + +def tiny_chatglm(legacy_max_length=None): + """Build a small CPU model without relying on Transformers' generation defaults.""" + config = ChatGLMConfig( + num_layers=1, + padded_vocab_size=64, + hidden_size=32, + ffn_hidden_size=64, + kv_channels=8, + num_attention_heads=4, + seq_length=32, + use_cache=False, + ) + # Transformers 4 supplies this generation field; Transformers 5 does not. + # Explicitly remove it so the fallback is covered on either version. + if hasattr(config, "max_length"): + del config.max_length + if legacy_max_length is not None: + config.max_length = legacy_max_length + model = ChatGLMForConditionalGeneration(config, empty_init=False, device="cpu") + for module in model.modules(): + if module.__class__.__name__ == "RMSNorm": + torch.nn.init.ones_(module.weight) + return model + + +def chatglm_inputs(): + """Return deterministic CPU input tokens and next-token labels.""" + tokens = torch.tensor([[1, 2, 3, 4]]) + return {"input_ids": tokens, "labels": tokens.clone()} + + +@pytest.mark.parametrize("use_cache", [False, True]) +def test_chatglm_config_retains_use_cache(use_cache): + config = ChatGLMConfig(use_cache=use_cache) + assert config.use_cache is use_cache + assert ChatGLMConfig.from_dict(config.to_dict()).use_cache is use_cache + + +@pytest.mark.parametrize("legacy_max_length", [None, 17]) +def test_chatglm_sequence_length_and_training_cpu(legacy_max_length): + model = tiny_chatglm(legacy_max_length) + assert model.max_sequence_length == (32 if legacy_max_length is None else legacy_max_length) + output = model(**chatglm_inputs()) + assert output.logits.shape == (1, 4, 64) + assert torch.isfinite(output.loss) + output.loss.backward() + gradients = [parameter.grad for parameter in model.parameters() if parameter.requires_grad] + assert all(gradient is not None and torch.isfinite(gradient).all() for gradient in gradients) + assert any(torch.count_nonzero(gradient).item() for gradient in gradients) + + +def test_chatglm_lazy_materialization_cpu(): + context = LazyInitContext(default_device="cpu") + with context: + model = tiny_chatglm() + assert all(isinstance(parameter, LazyTensor) for parameter in model.parameters()) + context.materialize(model) + assert model.max_sequence_length == 32 + assert all( + not isinstance(parameter, LazyTensor) and parameter.device.type == "cpu" for parameter in model.parameters() + ) + output = model(**chatglm_inputs()) + assert torch.isfinite(output.loss) + output.loss.backward() + assert all(parameter.grad is not None and torch.isfinite(parameter.grad).all() for parameter in model.parameters())