Skip to content

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

Closed
timmoon10 wants to merge 114 commits intoNVIDIA:mainfrom
timmoon10:fp8_cuda_graphs_transpose_caching
Closed

[PyTorch] Fix Float8Tensor transpose caching in #575#735
timmoon10 wants to merge 114 commits intoNVIDIA:mainfrom
timmoon10:fp8_cuda_graphs_transpose_caching

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

Signed-off-by: Tim Moon <tmoon@nvidia.com>
@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

timmoon10 commented Mar 27, 2024

/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