Skip to content

Pull requests: apple/axlearn

New pull request New
Clear current search query, filters, and sorts
0 Open 24 Closed
0 Open 24 Closed
Author
Filter by author
Loading
Label
Filter by label
Loading
Use alt + click/return to exclude labels
or ⇧ + click/return for logical OR
Projects
Filter by project
Loading
Milestones
Filter by milestone
Loading
Reviews
Filter by reviews
No reviews Review required Approved review Changes requested
Assignee
Filter by who’s assigned
Sort
Sort by
Newest Oldest Most commented Least commented Recently updated Least recently updated Best match
Most reactions
👍 👎 😄 🎉 😕 ❤️ 🚀 👀

Pull requests list

[Bug Fix] Fix a nit issue for Jax rollback
#856 by kelvin-zou was merged Nov 25, 2024 Loading…
[Bug fix] Jax rollback for TPU
#855 by kelvin-zou was closed Nov 23, 2024 Loading…
Rollback to Jax 0.4.33
#854 by kelvin-zou was merged Nov 22, 2024 Loading…
1
Speed up FA Backward pass in GPU via parallelizing sequence dimension
#818 by kelvin-zou was merged Nov 12, 2024 Loading…
6
[Bug fix] Update gpu flash attention after syntax change and fixed unit tests for flash attention
#809 by kelvin-zou was merged Nov 5, 2024 Loading…
@markblee
Add crash when suspecting a hang
#806 by kelvin-zou was merged Nov 4, 2024 Loading…
Bump up circle CI container size for build-and-test.
#794 by kelvin-zou was merged Oct 29, 2024 Loading…
[DO NOT CHECK IN, TEST ONLY]Test CI OOM issue
#792 by kelvin-zou was closed Oct 29, 2024 Loading…
Revert "Speed up Axlearn CI"
#790 by kelvin-zou was closed Oct 29, 2024 Loading…
2
Add TPU Monitoring for faster hang detection
#786 by kelvin-zou was merged Oct 29, 2024 Loading…
6
[Bug fix] better GPU flash attention compatibility
#744 by kelvin-zou was merged Oct 14, 2024 Loading…
6
[Bug fix] fix on the wrong env setup
#723 by kelvin-zou was merged Oct 2, 2024 Loading…
Add CuDNN fused MHA kernel to axlearn
#705 by kelvin-zou was merged Sep 19, 2024 Loading…
18
Support customized mesh rules to support different HWs
#696 by kelvin-zou was merged Sep 13, 2024 Loading…
39
Add litepod config for 70B model
#656 by kelvin-zou was closed Sep 5, 2024 • Draft
6
Jax upgrade 4 30
#653 by kelvin-zou was merged Sep 5, 2024 Loading…
2
Fix a missing check on FFN remat spec
#508 by kelvin-zou was merged Jun 5, 2024 Loading…
add flash attention to fuji
#506 by kelvin-zou was merged Jun 4, 2024 Loading…
add 70B model to the fuji model lib
#505 by kelvin-zou was merged Jun 4, 2024 Loading…
Te flash attention
#464 by kelvin-zou was closed Sep 18, 2024 • Draft
Perf optimize with unroll=8
#461 by kelvin-zou was closed Jun 4, 2024 Loading…
1
Triton flash attention
#454 by kelvin-zou was merged May 12, 2024 Loading…
5
Upgrade jax to 0.4.25
#434 by kelvin-zou was merged May 1, 2024 Loading…
Upgrade Jax to 0.4.25
#413 by kelvin-zou was merged Apr 23, 2024 Loading…
ProTip! Type g p on any issue or pull request to go back to the pull request listing page.