lightning-thunder icon indicating copy to clipboard operation
lightning-thunder copied to clipboard

Adds SymTypes to tree_flatten

Open kiya00 opened this issue 7 months ago • 1 comments

Before submitting
  • [ ] Was this discussed/approved via a Github issue? (no need for typos and docs improvements)
  • [ ] Did you read the contributor guideline, Pull Request section?
  • [ ] Did you make sure to update the docs?
  • [ ] Did you write any new necessary tests?

What does this PR do?

Fixes #2085 .

PR review

Anyone in the community is free to review the PR once the tests have passed. If we didn't discuss your PR in Github issues there's a high chance it will not be merged.

Did you have fun?

Make sure you had fun coding 🙃

kiya00 avatar May 20 '25 10:05 kiya00

Hi @t-vi , although I couldn't reproduce the error (I pulled the litgpt main and the thunder main and ran the test), but I believe this should fix the litgpt CI

root@d642a7329c41:/wayan/tmp1/litgpt# pip list|grep litgpt
litgpt                 0.5.9.dev1                    /wayan/tmp1/litgpt
root@d642a7329c41:/wayan/tmp1/litgpt# pytest tests/ext_thunder/test_thunder_networks.py -k test_hf_for_nemo[Qwen/Qwen2.5-7B-Instruct] -vs
========================================== test session starts ==========================================
platform linux -- Python 3.12.3, pytest-8.3.5, pluggy-1.6.0 -- /usr/bin/python3
cachedir: .pytest_cache
Test order randomisation NOT enabled. Enable with --random-order or --random-order-bucket=<bucket_type>
benchmark: 5.1.0 (defaults: timer=time.perf_counter disable_gc=False min_rounds=5 min_time=0.000005 max_time=1.0 calibration_precision=10 warmup=False warmup_iterations=100000)
hypothesis profile 'default'
rootdir: /wayan/tmp1/litgpt
configfile: pyproject.toml
plugins: timestamper-0.0.10, random-order-1.1.1, cov-6.1.1, timeout-2.3.1, xdist-3.6.1, benchmark-5.1.0, hypothesis-6.131.18, shard-0.1.2, mpi-0.6
collected 41 items / 40 deselected / 1 selected
Running 1 items in this shard: tests/ext_thunder/test_thunder_networks.py::test_hf_for_nemo[Qwen/Qwen2.5-7B-Instruct]

[2025-05-20 04:54:31] tests/ext_thunder/test_thunder_networks.py::test_hf_for_nemo[Qwen/Qwen2.5-7B-Instruct] Sliding Window Attention is enabled but not implemented for `sdpa`; unexpected results may be encountered.
PASSED

============================ 1 passed, 40 deselected, 19 warnings in 12.30s =============================

kiya00 avatar May 20 '25 11:05 kiya00

@t-vi, let's merge this to fix https://github.com/Lightning-AI/lightning-thunder/issues/2085.

IvanYashchuk avatar May 22 '25 11:05 IvanYashchuk