Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[PyTorch] Fix Float8Tensor transpose caching in #575 #735

Closed
wants to merge 114 commits into from
Closed

[PyTorch] Fix Float8Tensor transpose caching in #575 #735

wants to merge 114 commits into from

Conversation

timmoon10
Copy link
Collaborator

@timmoon10 timmoon10 commented Mar 27, 2024 •
edited
Loading

In a discussion with @ksivaman, we wanted the transpose function in Float8Tensor to have the following properties:

  • The cache is updated in-place
  • The transpose kernel can be passed a flag that causes a noop
  • The cache is returned by default if available
  • The cache is invalidated by in-place tensor ops
  • Its name hints that it is an internal, expert-level function

This feature set is important for both #575 and #707.

Closed by #575.

ksivaman and others added 30 commits February 2, 2024 06:21
@ksivaman
FP8 cuda graphs
8cb93ff 
Signed-off-by: Kirthi Shankar Sivamani <[email protected]>
@ksivaman
Merge branch 'main' into fp8_cuda_graphs
37ef2f7 
Signed-off-by: Kirthi Shankar Sivamani <[email protected]>
@ksivaman
Fix FP8 convergence
1d220aa 
Signed-off-by: Kirthi Shankar Sivamani <[email protected]>
@ksivaman
return non-None for ONNX
a9314eb 
Signed-off-by: Kirthi Shankar Sivamani <[email protected]>
@ksivaman
Merge branch 'main' into fp8_cuda_graphs
7a68197
@ksivaman
[WIP] static memory amax reduction
a7e539c 
Signed-off-by: Kirthi Shankar Sivamani <[email protected]>
@ksivaman
[WIP] cleanup
ddeb54d 
Signed-off-by: Kirthi Shankar Sivamani <[email protected]>
@ksivaman
Refine
98b3669 
Signed-off-by: Kirthi Shankar Sivamani <[email protected]>
@ksivaman
Merge branch 'main' into fp8_cuda_graphs
8c05d70
@ksivaman
Fix numerics with graph capture
9e379a7 
Move backward amax reduction outside modules
Fix amax addition to global buffer

Signed-off-by: Kirthi Shankar Sivamani <[email protected]>
@ksivaman
Hook fixes
0d2a4a6 
Signed-off-by: Kirthi Shankar Sivamani <[email protected]>
@ksivaman
Cleanup
409f601 
Signed-off-by: Kirthi Shankar Sivamani <[email protected]>
@ksivaman
merge fused amax and scale update kernel
7b29e96 
Signed-off-by: Kirthi Shankar Sivamani <[email protected]>
@ksivaman
simple fusion
9784c0d 
Signed-off-by: Kirthi Shankar Sivamani <[email protected]>
@ksivaman
Skip fwd amax reduction during graph capture
3455c80 
Signed-off-by: Kirthi Shankar Sivamani <[email protected]>
@ksivaman
noop c+t kernel
ab26eb6 
Signed-off-by: Kirthi Shankar Sivamani <[email protected]>
@ksivaman
fix
9506b7e 
Signed-off-by: Kirthi Shankar Sivamani <[email protected]>
@ksivaman
Add caching
5952f56 
Signed-off-by: Kirthi Shankar Sivamani <[email protected]>
@ksivaman
Use outer (user) FP8 autocast to determine freq of bwd amax reduction
374867a 
Signed-off-by: Kirthi Shankar Sivamani <[email protected]>
@ksivaman
Merge branch 'fp8_cuda_graphs' into fp8_cuda_graphs_with_caching
ff2a8ff
@ksivaman
Compile
75978b0 
Signed-off-by: Kirthi Shankar Sivamani <[email protected]>
@ksivaman
fix graph case
ecd80dd 
Signed-off-by: Kirthi Shankar Sivamani <[email protected]>
@ksivaman
Fix
4230442 
Signed-off-by: Kirthi Shankar Sivamani <[email protected]>
@ksivaman
Merge branch 'main' into fp8_cuda_graphs
50b7d95
@ksivaman
remove alloc
11c48ed 
Signed-off-by: Kirthi Shankar Sivamani <[email protected]>
@ksivaman
Merge branch 'fp8_cuda_graphs' into fp8_cuda_graphs_with_caching
949f55a
@ksivaman
Working
55e1c7f 
Signed-off-by: Kirthi Shankar Sivamani <[email protected]>
@cyanguwa
add fused kernel for bulk update of amax and scales after reduction
b9c954a 
Signed-off-by: Charlene Yang <[email protected]>
@ksivaman
Merge branch 'main' into fp8_cuda_graphs
46b3c34
@cyanguwa
calculate a more accurate param limit
23222c7 
Signed-off-by: Charlene Yang <[email protected]>
ksivaman and others added 22 commits March 21, 2024 21:42
@ksivaman
meta device and numerics tests fixes
c387ff0 
Signed-off-by: Kirthi Shankar Sivamani <[email protected]>
@ksivaman
fix fused attention
5d29755 
Signed-off-by: Kirthi Shankar Sivamani <[email protected]>
@ksivaman
resolve conflicts
c54fbdb 
Signed-off-by: Kirthi Shankar Sivamani <[email protected]>
@ksivaman
Better design for fp8 weight caching
3e51aee 
Signed-off-by: Kirthi Shankar Sivamani <[email protected]>
@ksivaman
Remove testing stuff
31c7888 
Signed-off-by: Kirthi Shankar Sivamani <[email protected]>
@ksivaman
Float8Tensor transpose change API
aca4211 
Signed-off-by: Kirthi Shankar Sivamani <[email protected]>
@ksivaman
Improvements and review
f3c377f 
Signed-off-by: Kirthi Shankar Sivamani <[email protected]>
@ksivaman
fix dynamic amax history
bb5b4d6 
Signed-off-by: Kirthi Shankar Sivamani <[email protected]>
@ksivaman
Merge branch 'main' into fp8_cuda_graphs
324360b
@ksivaman
fix
6dbf7b3 
Signed-off-by: Kirthi Shankar Sivamani <[email protected]>
@ksivaman
Re-add kernel for paddle
5cb6eed 
Signed-off-by: Kirthi Shankar Sivamani <[email protected]>
@ksivaman
fix
0e56285 
Signed-off-by: Kirthi Shankar Sivamani <[email protected]>
@ksivaman
fix float8tensor test
1cd9438 
Signed-off-by: Kirthi Shankar Sivamani <[email protected]>
@ksivaman
fix float8tensor test
77a34e3 
Signed-off-by: Kirthi Shankar Sivamani <[email protected]>
@ksivaman
fix
1b87616 
Signed-off-by: Kirthi Shankar Sivamani <[email protected]>
@ksivaman
fix
07f262e 
Signed-off-by: Kirthi Shankar Sivamani <[email protected]>
@ksivaman
Fix
4bcbb66 
Signed-off-by: Kirthi Shankar Sivamani <[email protected]>
@ksivaman
Fix s_inv compute
48f7005 
Signed-off-by: Kirthi Shankar Sivamani <[email protected]>
@ksivaman
Add additional checks
49a7964 
Signed-off-by: Kirthi Shankar Sivamani <[email protected]>
@ksivaman
Cache norm workspace/barriers/partial grads
404f8fa 
Signed-off-by: Kirthi Shankar Sivamani <[email protected]>
@ksivaman
bug fixes
51e8f64 
Signed-off-by: Kirthi Shankar Sivamani <[email protected]>
@timmoon10
Use lazy caching by default for Float8Tensor transpose
fa9a627 
Restore general support for in-place operations.

Signed-off-by: Tim Moon <[email protected]>
@timmoon10 timmoon10 requested a review from ksivaman March 27, 2024 08:27
@timmoon10 timmoon10 marked this pull request as draft March 27, 2024 08:27
@timmoon10
Copy link
Collaborator Author

/te-ci pytorch

@timmoon10 timmoon10 mentioned this pull request Mar 27, 2024
Merged
@ksivaman
Copy link
Member

ksivaman commented Apr 3, 2024

@timmoon10 Closing

@ksivaman ksivaman closed this Apr 3, 2024
@timmoon10 timmoon10 mentioned this pull request Apr 5, 2024
Merged
8 tasks
@timmoon10 timmoon10 mentioned this pull request May 16, 2024
Merged
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@ksivaman ksivaman Awaiting requested review from ksivaman

Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
3 participants
@timmoon10 @ksivaman @cyanguwa