Fix deepspeed ci (#48640)
* Rebuild the DeepSpeed CI image on a plain CUDA 12.6 base
The image was pinned to `PYTORCH=2.8.0` while torchaudio moved on. torchaudio is
in maintenance mode and its final 2.11.0 release no longer pins `torch`, so pip
resolved `torchaudio-2.11.0+cu126` next to `torch-2.8.0+cu126` and the mismatched
extension failed to load:
OSError: .../torchaudio/lib/_torchaudio.abi3.so: undefined symbol: torch_library_impl
That breaks `import transformers` through `audio_utils`, and DeepSpeed imports
transformers while generating its metadata, so the build ended in
`metadata-generation-failed`. The build has been red daily since at least
2026-09-04, and Daily CI reinstalls DeepSpeed inside the GPU VM where it hit the
same error during Setup -- which skipped every downstream DeepSpeed job,
`run_trainer_and_fsdp_gpu` included.
Bumping torch to the version `transformers-all-latest-gpu` uses puts it inside
torchaudio 2.11's supported range and needs no extra pinning. Move off
`nvcr.io/nvidia/pytorch:24.08-py3` onto a plain `nvidia/cuda` base like that image
uses, which drops everything that only existed to undo what the nvcr image
shipped: the `pandas<2`/`numpy<2` pin, uninstalling `transformer-engine`,
`torch-tensorrt`, `nvtx` and `apex`, and repairing `pydantic`. It also drops the
`hypothesis==5.35.1` that image preinstalled, which pytest could no longer load
once torch 2.13 pulled in a setuptools without `pkg_resources`.
Unlike the general test image this one stays on CUDA 12.x: DeepSpeed compiles its
own ops, and `get_default_compute_capabilities()` only knows about CUDA 11 and 12,
so on CUDA 13 it cross-compiles for `6.0;6.1;7.0` and nvcc rejects the lot. The
Dockerfile records this so the base can be moved forward once DeepSpeed catches
up. `libaio-dev`, `WORKDIR /workspace` and a narrowed `TORCH_CUDA_ARCH_LIST` are
kept on top of the general image's setup, all commented in place.
Also pull in the `sklearn` extra: the model zoo tests run the example scripts, which
compute metrics with scikit-learn and scipy. The general test image gets those via
`[dev]`, and the nvcr base used to preinstall them.
Also install `python-is-python3`: the model zoo tests execute DeepSpeed`s
`zero_to_fp32.py` directly, and its `#!/usr/bin/env python` shebang found nothing on
an Ubuntu base that only ships `python3`, so all 54 of those tests exited 127.
* Fix DeepSpeed Trainer tests against deepspeed>=0.19.6
`deepspeed>=0.19.6` rejects a non-positive `warmup_num_steps` in `WarmupLR`,
`WarmupDecayLR` and `WarmupCosineLR`, while older versions silently clamped it
to `max(2, warmup_num_steps)`. `trainer_config_finalize` fills the DeepSpeed
scheduler config with `args.get_warmup_steps()`, which is `0` for the default
`warmup_steps=0`/`warmup_ratio=0.0`, so every DeepSpeed Trainer run died with
`ValueError: warmup_num_steps must be a positive integer, got 0`.
Keep the auto-filled value positive. `1` is clamped to `2` by DeepSpeed just as
`0` used to be, so no-warmup runs behave exactly as before.
Also give `test_auto_batch_size_with_deepspeed` a rank environment. It builds a
`Trainer` in-process rather than under a launcher, and DeepSpeed refuses to
initialize without `LOCAL_RANK`, so it failed with an assertion from
`_do_args_sanity_check`. Every other in-process DeepSpeed Trainer test already
wraps itself in `mockenv_context(**dist_env_1_gpu)`.
* Give the grad accumulation grad-norm check its own tolerance
`test_gradient_accumulation_grad_norm_without_num_items_in_batch` compares grad
norms between a `batch_size=8, gas=1` baseline and an equivalent
`batch_size=4, gas=2` run. With `model_accepts_loss_kwargs=False` there is no
`num_items_in_batch`, so every micro-batch is mean-reduced over its own valid
label count -- and the fixture pads to `max_length=16` with pad labels masked to
`-100`, so those counts genuinely differ between samples. Regrouping the same
samples into smaller micro-batches therefore shifts both the loss and the grad
norm, which is why the loss already gets a looser tolerance here.
`DataParallel` splits every micro-batch again across replicas (8 -> 4+4 and
4 -> 2+2), regrouping them more finely still and pushing the grad norm ratio to
~1.11 against the hardcoded 0.1 delta. Give the grad-norm check the same
per-caller tolerance the loss check has and loosen it for this case only. A real
GAS leak shows a ratio near `gas_steps`, so 0.2 still catches it, and the other
callers keep the tight 0.1.
* Give the distributed test collator position_ids for Ulysses SP
`deepspeed>=0.19.6` requires `position_ids` in every batch under Ulysses sequence
parallelism, so that a token keeps its correct global position once the sequence
is sharded across ranks:
ValueError: Ulysses SP requires `position_ids` in every dataloader batch ...
Ensure your data collator includes position_ids in its output.
`DataCollatorForLanguageModeling` emits `input_ids`, `attention_mask` and
`labels` only, so `test_alst_ulysses_sp` died in `ulysses_sp.py`'s `refill()`
before training started. Wrap the collator and add `torch.arange(seq_len)` per
sample, which is what DeepSpeed documents for non-packed sequences and also what
the model derives internally when `position_ids` is absent -- so the non-SP half
of the comparison this test makes is unaffected.
