DeepSpeed

2d80cedd - Cap CPU accelerator memory at the cgroup limit of the process (#8930)

Cap CPU accelerator memory at the cgroup limit of the process (#8930) **Issue** `CPU_Accelerator.total_memory()` and `available_memory()` return host RAM from `psutil.virtual_memory()`. In a container or a Kubernetes pod, the process can use only its cgroup limit. Callers that plan memory from these values on the CPU accelerator include ZeRO 1/2 flattening (`stage_1_and_2.py:451`) and the autotuner. Fixes #8929. Related: #1891 asks for cgroup-aware CPU memory limits. #3458 (Docker), #2632 and #2788 (Slurm) report processes killed by a cgroup limit without a message. **Fix** Only the CPU accelerator, its tests and its docs change. - `total_memory()` returns the smallest memory limit from the process cgroup up to the root, capped at host RAM. - `available_memory()` returns the smallest free headroom over every limited cgroup on that path. A parent limit also counts the usage of its other children, so each limited level is checked. Usage excludes the inactive page cache, which the kernel can reclaim. - The cgroup mount point and mount root come from `/proc/self/mountinfo`. When several mounts match, the one whose root contains the process cgroup is used. This supports mounts that expose a subtree, cgroup v1 co-mounts such as `cpu,memory`, and delegated subtrees mounted before the main mount. - On by default, as asked in review. `RESPECT_CGROUP_MEMORY_BY_DEFAULT = True` is the one line that sets the default. `DS_CPU_RESPECT_CGROUP_MEMORY=0` turns it off, `=1` turns it on. `export_envs()` returns the variable, so the launcher forwards it to every node. - The accelerator setup guide documents the variable. `abstract_accelerator.py` does not change. **How to test** Needs Docker with cgroup v2. Save the first script from the collapsed sections at the end as `check_cgroup_memory.sh`. One command installs `master` (before) and this branch (after) in the same container with 1 GiB. It prints `total_memory()` and `available_memory()` from three cgroups with `DS_CPU_RESPECT_CGROUP_MEMORY` unset, and once more with `=0`. ```sh sh check_cgroup_memory.sh "git+https://github.com/deepspeedai/DeepSpeed@master" \ "git+https://github.com/prabhkodes/DeepSpeed@cpu-accelerator-cgroup-memory" ``` Output (Ubuntu 24.04 VM, kernel 6.8, 3.8 GiB, cgroup v2, Docker 29.1 with the systemd driver): ``` process cgroup (container limit 1 GiB) total available before A. container root 3904 MiB 2773 MiB before B. child cgroup, limit 256 MiB 3904 MiB 2774 MiB before C. child 600 MiB; parent 800 MiB, sibling uses 500 3904 MiB 2685 MiB after A. container root 1024 MiB 342 MiB after B. child cgroup, limit 256 MiB 256 MiB 90 MiB after C. child 600 MiB; parent 800 MiB, sibling uses 500 600 MiB 129 MiB after A. container root, DS_CPU_RESPECT_CGROUP_MEMORY=0 3904 MiB 2755 MiB ``` Before, every case reports the host. After, `total` is the limit of each case. In C, `available` is the free headroom of the 800 MiB parent, which the 500 MiB sibling also uses. In A, the container's 1 GiB also holds that sibling. With `=0`, the values are host RAM again. On a cgroup v1 host, the second script, `check_cgroup_v1.sh`, runs the same cases with the same arguments. Same VM booted with `systemd.unified_cgroup_hierarchy=0`, Docker 29.1 with the cgroupfs driver and its default host cgroup namespace: ``` process cgroup (container limit 1 GiB) total available before A. container root 3904 MiB 2850 MiB before B. child cgroup, limit 256 MiB 3904 MiB 2814 MiB before C. child 600 MiB; parent 800 MiB, sibling uses 500 3904 MiB 2781 MiB after A. container root 1024 MiB 175 MiB after B. child cgroup, limit 256 MiB 256 MiB 89 MiB after C. child 600 MiB; parent 800 MiB, sibling uses 500 600 MiB 129 MiB after A. container root, DS_CPU_RESPECT_CGROUP_MEMORY=0 3904 MiB 2672 MiB ``` Each script ran twice per cgroup version. `total` and the C row are identical in all runs. `available` in A and B reads live usage, so it moves between runs. **Proof** ZeRO training on the CPU accelerator, three Adam steps, flag `0` and `1`, two fresh containers each: | processes | precision | ZeRO 1 | ZeRO 2 | ZeRO 3 | |---|---|---|---|---| | 1 | fp32 | identical losses | identical losses | identical losses | | 1 | bf16 | identical losses | identical losses | identical losses | | 2 (`deepspeed` launcher, gloo) | fp32 | identical losses | identical losses | identical losses | **Performance** Throughput and peak RSS, flag `0` and `1` (`1` is the code path that is now the default), 30 timed steps on a 4×1024 MLP. Four runs per cell, two fresh containers, flag alternated. Median steps/s with range: | processes | precision | ZeRO | steps/s, flag 0 | steps/s, flag 1 | peak RSS, flag 0 / 1 | |---|---|---|---|---|---| | 1 | fp32 | 1 | 23.8 (22.9–24.7) | 24.2 (23.2–25.2) | 575 / 575 MiB | | 1 | fp32 | 2 | 24.3 (24.0–24.9) | 23.3 (22.4–24.1) | 572 / 571 MiB | | 1 | fp32 | 3 | 22.1 (21.1–26.5) | 22.1 (20.9–22.7) | 606 / 606 MiB | | 1 | bf16 | 1 | 41.5 (37.7–42.4) | 41.2 (40.7–42.5) | 530 / 530 MiB | | 1 | bf16 | 2 | 42.0 (40.1–42.5) | 41.6 (33.5–42.6) | 530 / 530 MiB | | 1 | bf16 | 3 | 54.1 (45.2–55.4) | 54.2 (51.3–56.7) | 556 / 556 MiB | | 2 | fp32 | 1 | 11.1 (10.8–11.4) | 11.4 (10.4–11.7) | 524 / 524 MiB | | 2 | fp32 | 2 | 11.2 (10.3–12.4) | 11.9 (11.2–12.6) | 524 / 522 MiB | | 2 | fp32 | 3 | 11.0 (10.7–11.1) | 10.6 (9.9–11.0) | 564 / 564 MiB | Peak RSS does not change. The throughput differences go in both directions and stay inside the run-to-run range. A call counter shows why: ZeRO 1 and 2 call the memory queries once at setup, ZeRO 3 does not call them, and no configuration calls them during training steps. Cost of one query, best of 5 × 5000 calls, process three levels below the cgroup root: | | `total_memory()` | `available_memory()` | |---|---|---| | master | 12.4 µs | 12.5 µs | | this PR, `=0` | 13.1 µs | 12.9 µs | | this PR, default (on) | 81.5 µs | 80.7 µs | Measured on ca6a5a6, before the default changed; the two rows are the same code paths. With the cap on, a query reads `/proc/self/cgroup`, `/proc/self/mountinfo` and the cgroup files: about 70 µs more, once per run. ``` pytest tests/unit/v1/accelerator/test_accelerator.py # cgroup v1 and v2 host: 33 passed, 2 skipped # in a 1 GiB container: 30 passed, 5 skipped (the pin_memory op cannot be built there) yapf 0.43.0, flake8 5.0.4 --config=.flake8 # the pre-commit versions: no changes, no findings ``` Twelve new tests. They cover a limit on a parent cgroup, the smallest limit, the v1 sentinel, the headroom of every limited level, mountinfo resolution with several mounts, the launcher export, and default × override. One test uses `/proc/self/cgroup`, the cgroup mountinfo lines and the memory files captured in a Docker container on the cgroup v1 host above, so the v1 path is checked against a real layout. With the mount-root strip removed from the code, that test fails and the hand-written v1 tests still pass. **Hardware.** Apple M4 host, Docker Desktop Linux 6.10 VM (aarch64, 10 CPUs, 7.8 GiB, cgroup v2). Python 3.12.16, torch 2.14.1+cpu, `DS_BUILD_OPS=0`. No GPU. cgroup v1, the v2 re-run and the unit tests above: Lima VM on the same host, Ubuntu 24.04, kernel 6.8.0-142 (aarch64, 4 vCPU, 3.8 GiB), Docker 29.1.3, Python 3.12.16, torch 2.15.0+cpu. **Compatibility matrix.** - Validated: CPU accelerator; ZeRO 1, 2 and 3; fp32 and bf16; 1 and 2 processes; no offload. Memory values on cgroup v1 and v2. - Not reachable on the CPU accelerator: `deepspeed.compile` (`is_deepcompile_supported()` requires CUDA) and the inference v2 KV cache (CUDA kernels only). - Not validated: CPU and NVMe offload, fp16, multi-node runs (the launcher export is unit-tested only). - Not affected: CUDA, XPU, NPU, MLU, HPU, SDAA, SUPA and MPS read their own device memory. `see_memory_usage()` and the timer read host RAM only to log it. No open PR changes the CPU accelerator memory functions (searched `cgroup`, `cpu_accelerator available_memory`, `cpu accelerator memory`, `1891 in:body` on 2026-10-27). AI assistance was used: an AI coding agent (Claude Code) wrote the change and the tests; I reviewed each changed line and ran the tests, the reproduction and the ZeRO runs. <details> <summary>check_cgroup_memory.sh (cgroup v2)</summary> ```sh #!/bin/sh # Before/after check of DeepSpeed's CPU accelerator memory values inside cgroups. Needs Docker with cgroup v2. # sh check_cgroup_memory.sh <before pip spec> <after pip spec> # Example: # sh check_cgroup_memory.sh "git+https://github.com/deepspeedai/DeepSpeed@master" \ # "git+https://github.com/prabhkodes/DeepSpeed@cpu-accelerator-cgroup-memory" # A pip spec can also be a source tarball in the directory given as SRC_DIR (mounted at /src). # IMG replaces the python:3.12-slim image, e.g. with one that has torch installed. set -eu docker run --rm -i --privileged --cgroupns=private --memory=1g ${EXTRA_DOCKER_ARGS:-} ${SRC_DIR:+-v "$SRC_DIR:/src:ro"} \ -e BEFORE="${1:?before spec}" -e AFTER="${2:?after spec}" "${IMG:-python:3.12-slim}" sh -s <<'INNER' set -e export DEBIAN_FRONTEND=noninteractive apt-get -qq update >/dev/null 2>&1 && apt-get -qq install -y git >/dev/null 2>&1 pip install -q torch --index-url https://download.pytorch.org/whl/cpu >/dev/null 2>&1 pip install -q psutil numpy pydantic packaging py-cpuinfo tqdm hjson msgpack ninja einops >/dev/null 2>&1 cat > /tmp/probe.py <<'PY' from deepspeed.accelerator import get_accelerator a = get_accelerator() print(f"{a.total_memory() >> 20:>7} MiB {a.available_memory() >> 20:>7} MiB") PY cat > /tmp/hold.py <<'PY' import time b = b"\x01" * (500 << 20) time.sleep(3600) PY C=/sys/fs/cgroup mkdir -p $C/init; for p in $(cat $C/cgroup.procs); do echo $p > $C/init/cgroup.procs 2>/dev/null || true; done echo "+memory" > $C/cgroup.subtree_control mkdir -p $C/nested; echo $((256 << 20)) > $C/nested/memory.max mkdir -p $C/shared; echo $((800 << 20)) > $C/shared/memory.max; echo 0 > $C/shared/memory.swap.max echo "+memory" > $C/shared/cgroup.subtree_control mkdir -p $C/shared/sibling $C/shared/me; echo $((600 << 20)) > $C/shared/me/memory.max sh -c "echo \$\$ > $C/shared/sibling/cgroup.procs; echo -1000 > /proc/self/oom_score_adj; exec python /tmp/hold.py" & sleep 3 # run <version> <label> <cgroup> [VAR=value]: the probe runs with the flag unset unless given. run() { printf "%-7s %-56s" "$1" "$2" sh -c "echo \$\$ > $C/$3/cgroup.procs; DS_ACCELERATOR=cpu ${4:-} exec python /tmp/probe.py 2>/dev/null" } printf "%-7s %-56s %11s %11s\n" "" "process cgroup (container limit 1 GiB)" "total" "available" for v in before after; do if [ $v = before ]; then spec=$BEFORE; else spec=$AFTER; fi DS_BUILD_OPS=0 pip install -q --no-deps --force-reinstall "$spec" >/dev/null 2>&1 run $v "A. container root" init run $v "B. child cgroup, limit 256 MiB" nested run $v "C. child 600 MiB; parent 800 MiB, sibling uses 500" shared/me done run after "A. container root, DS_CPU_RESPECT_CGROUP_MEMORY=0" init DS_CPU_RESPECT_CGROUP_MEMORY=0 INNER ``` </details> <details> <summary>check_cgroup_v1.sh (cgroup v1)</summary> ```sh #!/bin/sh # Before/after check of DeepSpeed's CPU accelerator memory values on a cgroup v1 host. Run as root on the host. # Docker's default on cgroup v1 is the host cgroup namespace: /proc/self/cgroup shows the host path of the # container, and the cgroup mount inside the container starts at the container's own cgroup. # sh check_cgroup_v1.sh <before pip spec> <after pip spec> # Example: # sh check_cgroup_v1.sh "git+https://github.com/deepspeedai/DeepSpeed@master" \ # "git+https://github.com/prabhkodes/DeepSpeed@cpu-accelerator-cgroup-memory" # A pip spec can also be a source tarball in the directory given as SRC_DIR (mounted at /src). # IMG replaces the python:3.12-slim image, e.g. with one that has torch installed. set -eu [ "$(stat -fc %T /sys/fs/cgroup)" = tmpfs ] || { echo "needs a cgroup v1 host" >&2; exit 1; } docker run --rm -i --privileged --memory=1g ${SRC_DIR:+-v "$SRC_DIR:/src:ro"} \ -e BEFORE="${1:?before spec}" -e AFTER="${2:?after spec}" "${IMG:-python:3.12-slim}" sh -s <<'INNER' set -e export DEBIAN_FRONTEND=noninteractive apt-get -qq update >/dev/null 2>&1 && apt-get -qq install -y git >/dev/null 2>&1 pip install -q torch --index-url https://download.pytorch.org/whl/cpu >/dev/null 2>&1 pip install -q psutil numpy pydantic packaging py-cpuinfo tqdm hjson msgpack ninja einops >/dev/null 2>&1 cat > /tmp/probe.py <<'PY' from deepspeed.accelerator import get_accelerator a = get_accelerator() print(f"{a.total_memory() >> 20:>7} MiB {a.available_memory() >> 20:>7} MiB") PY cat > /tmp/hold.py <<'PY' import time b = b"\x01" * (500 << 20) time.sleep(3600) PY C=/sys/fs/cgroup/memory mkdir -p $C/nested; echo $((256 << 20)) > $C/nested/memory.limit_in_bytes mkdir -p $C/shared/sibling $C/shared/me echo $((800 << 20)) > $C/shared/memory.limit_in_bytes echo $((600 << 20)) > $C/shared/me/memory.limit_in_bytes sh -c "echo \$\$ > $C/shared/sibling/cgroup.procs; echo -1000 > /proc/self/oom_score_adj; exec python /tmp/hold.py" & sleep 3 # run <version> <label> <cgroup under the container> [VAR=value]: the probe runs with the flag unset unless given. run() { printf "%-7s %-56s" "$1" "$2" sh -c "echo \$\$ > $C/$3/cgroup.procs; DS_ACCELERATOR=cpu ${4:-} exec python /tmp/probe.py 2>/dev/null" } printf "%-7s %-56s %11s %11s\n" "" "process cgroup (container limit 1 GiB)" "total" "available" for v in before after; do if [ $v = before ]; then spec=$BEFORE; else spec=$AFTER; fi DS_BUILD_OPS=0 pip install -q --no-deps --force-reinstall "$spec" >/dev/null 2>&1 run $v "A. container root" . run $v "B. child cgroup, limit 256 MiB" nested run $v "C. child 600 MiB; parent 800 MiB, sibling uses 500" shared/me done run after "A. container root, DS_CPU_RESPECT_CGROUP_MEMORY=0" . DS_CPU_RESPECT_CGROUP_MEMORY=0 # The files the CPU accelerator reads, as seen from cgroup C. sh -c "echo \$\$ > $C/shared/me/cgroup.procs echo '### /proc/self/cgroup'; cat /proc/self/cgroup echo '### /proc/self/mountinfo (cgroup lines)'; grep -E ' - cgroup2? ' /proc/self/mountinfo for d in $C/shared/me $C/shared $C; do echo \"### \$d\" for f in memory.limit_in_bytes memory.usage_in_bytes; do echo \"\$f \$(cat \$d/\$f)\"; done grep -E '^total_inactive_file ' \$d/memory.stat done echo '### /sys/fs/cgroup/unified'; ls /sys/fs/cgroup/unified | tr '\n' ' '; echo" INNER ``` </details> **Precedent.** The same walk is already in production elsewhere. Rust's `std::thread::available_parallelism` takes the smallest `cpu.max` from the process cgroup up to the mount ([`quota_v2`](https://github.com/rust-lang/rust/blob/41089d07a0a2afc16adfe28a760d94d73d7298a4/library/std/src/sys/thread/unix.rs#L908-L953)). The JDK fixed limits on parent nested cgroups in [JDK-8322420](https://bugs.openjdk.org/browse/JDK-8322420) (2024) and the hierarchical memory limit in [JDK-8370572](https://bugs.openjdk.org/browse/JDK-8370572) (2025). Go 1.25 finds the process cgroup through `/proc/self/cgroup` and `/proc/self/mountinfo` ([`cgroup_linux.go`](https://github.com/golang/go/blob/63313a0fc5b68d86399a64e323113667d8aa7823/src/runtime/cgroup_linux.go#L11-L24)), reading the leaf only. vLLM's CPU-quota reader walks the hierarchy ([`_cgroup_cpu_limit`](https://github.com/vllm-project/vllm/blob/ed6b8f4cd2b47e4240e7ff5feb5105444245b0c2/vllm/utils/torch_utils.py#L180-L231)); vllm-project/vllm#64104 adds the same for memory. The same fix for OpenObserve merged on 2026-10-28 (openobserve/openobserve#15906). --------- Signed-off-by: Prabhsharan Singh <228956661+prabhkodes@users.noreply.github.com> Co-authored-by: Prabhsharan Singh <228956661+prabhkodes@users.noreply.github.com>