Skip to content

Reduce chunked cross-entropy memory test footprint - #8824

Merged
delock merged 2 commits into
deepspeedai:masterfrom
tohtana:tohtana/reduce-cross-entropy-test-memory
Oct 15, 2026
Merged

delock merged 2 commits into
deepspeedai:masterfrom
tohtana:tohtana/reduce-cross-entropy-test-memory

Conversation

@tohtana

@tohtana tohtana commented Oct 14, 2026 •

Copy link
Copy Markdown
Collaborator

The original 8,192-token by 151,936-vocabulary test needs 4.64 GiB for each full FP32 logits-sized tensor. In the reproduced CI stall, a long-lived pytest worker held 28,832 MiB (28.2 GiB) of GPU memory.

This PR reduces the token count from 8,192 to 128 while retaining the 151,936-entry vocabulary: each full FP32 tensor becomes 74.2 MiB, a 64-fold reduction. Preserve loss/gradient tolerances and the two-FP32-tensor memory-savings requirement, and release test-owned CUDA references and cache state on exit.

We now compare gradients in eight-row blocks and calculate the match fraction from integer counts instead of an FP32 mean. The 99.9% threshold is unchanged, but boundary rounding can differ. Check both torch (block_rows=8) and available Triton, instead of selecting just one backend.

Signed-off-by: Masahiro Tanaka <tanaka.masahiro@gmail.com>
@chatgpt-codex-connector

chatgpt-codex-connector Bot commented Oct 14, 2026 •

Copy link
Copy Markdown

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review Status Commit Review trigger
📝 Code Review ✅ Completed 2026-10-14T20:36:12.922133Z 0c446ac PR opened
ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

Signed-off-by: Masahiro Tanaka <tanaka.masahiro@gmail.com>
@delock
delock added this pull request to the merge queue Oct 15, 2026
Merged via the queue into deepspeedai:master with commit 5f266e2 Oct 15, 2026
15 checks passed
banxingmjj pushed a commit to openanolis/DeepSpeed that referenced this pull request Oct 15, 2026
(deepspeedai#8824 should be merged first)

In deepspeedai#8770, we pinned Modal sandbox placement to Oracle Cloud
Infrastructure (OCI), where a working NVIDIA driver version was still
available, to avoid the CUDA/NCCL failures seen on other Modal hosts.
The affected driver has since been rolled out to OCI as well, so this
workaround no longer provides a reliable path for our multi-GPU CI
workload. Modal has identified the issue and is working on a fix.

This PR temporarily disables Modal until the driver issue is resolved,
and run the required torch-latest workflow entirely on the existing AWS
backend.

---------

Signed-off-by: Masahiro Tanaka <tanaka.masahiro@gmail.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants