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 c977b552e2da..4f4b3909e8e6 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 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 e5b607315517..af44ef49ffe6 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.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 | | setuptools | 84.0.0 | | PyTorch's Python cuDNN package | nvidia-cudnn-cu13==9.20.0.48 | @@ -53,12 +55,70 @@ 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.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. 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. +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 \ + 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 \ + tests/test_shardformer/test_hf5_compat/test_integration_loading.py \ + tests/test_shardformer/test_hf5_compat/test_qwen_loss.py +``` + +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. + ## 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..983f54fbf56f 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.21.2 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/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 01ed90e15fd8..fe3a06f574f7 100644 --- a/docker/torch213/verify_image.py +++ b/docker/torch213/verify_image.py @@ -6,12 +6,142 @@ 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, checkpoint reload and resumed training.""" + 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, + ) + base_model = LlamaForCausalLM(config).to(device="cpu", dtype=torch.float32) + model = get_peft_model( + deepcopy(base_model), + 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" + 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) + + +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 +183,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..fc6b13e2c1ef 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.21.2 bitsandbytes>=0.39.0 rpyc==6.0.0 fastapi uvicorn==0.29.0 galore_torch -diffusers==0.29.0 +diffusers==0.35.2 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())