Lightning-AI / Lightning-AI/lightning-thunder

reenable apex backward / phantom grad tests without wrecking CI

Open
#392 3 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

apex autograd ci / tests memory use
Dominant language
Python
Stars
1.5k
Forks
121
PR merge metrics
No merged PRs in 30d

Description

I'm disabling the apex tests because they use excessively much memory.
Not sure whether this is the tests or a bug in the executor.
I'd be happy to see these re-enabled either with lower memory consumption or executed without parallelism.

test cuda memory use test_apex_cross_entropy_backward[cuda-float16] 4 memory 1.721745491027832
test cuda memory use test_apex_cross_entropy_backward[cuda-bfloat16] 5 memory 1.721745491027832
test cuda memory use test_apex_cross_entropy_backward[cuda-float32] 6 memory 2.872036933898926
test cuda memory use test_apex_cross_entropy_phantom_grad[cuda-float16] 7 memory 8.822380542755127
test cuda memory use test_apex_cross_entropy_phantom_grad[cuda-bfloat16] 8 memory 9.014111042022705
test cuda memory use test_apex_cross_entropy_phantom_grad[cuda-float32] 9 memory 9.401249408721924

cc @borda @crcrpar

Contributor guide

No contributing guide indexed for this repository

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start with the named test_apex_cross_entropy_backward and test_apex_cross_entropy_phantom_grad cases and the Apex test configuration that currently disables them. Measure their CUDA memory use under parallel and non-parallel execution, then verify the tests can be re-enabled without exhausting CI memory.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
ci-cd, performance, testing-qa
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.