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

[FSDP2][1/n] construct NF4Tensor from bf16/fp16/fp32 #118

Merged
merged 22 commits into from
Apr 19, 2024
Merged

[FSDP2][1/n] construct NF4Tensor from bf16/fp16/fp32 #118

merged 22 commits into from
Apr 19, 2024

Conversation

weifengpy
Copy link
Contributor

@weifengpy weifengpy commented Apr 4, 2024 •
edited
Loading

  • addressing to(bf16/fp16/fp32) with __torch_function__, as suggested by @cpuhrsch
  • NF4Tensor.dtype is the dtype from construction. For example, to_nf4(fp32_tensor) will return nf4tensor with dtype fp32
  • dtype=fp16/bf16/fp32 in pytest test/dtypes/test_nf4.py
  • tested in torchtune with pytest tests -m integration_test @rohan-varma

it brings 2 benefits

@weifengpy
proof of concept for FSDP2 + NF4Tensor
0a13e6a 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
@facebook-github-bot facebook-github-bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Apr 4, 2024
@weifengpy weifengpy marked this pull request as draft April 4, 2024 01:21
@weifengpy weifengpy mentioned this pull request Apr 4, 2024
Closed
@weifengpy weifengpy changed the title proof of concept for FSDP2 + NF4Tensor [In Progress] FSDP2 + NF4Tensor Apr 4, 2024
cpuhrsch and others added 3 commits April 4, 2024 10:53
@cpuhrsch
Merge branch 'main' into main
9a56eaa
@weifengpy
fsdp extention for tensor subclass
8180540 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
@weifengpy
support fp32
95b03e1 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
@weifengpy

This comment was marked as outdated.

Sign in to view
@weifengpy weifengpy mentioned this pull request Apr 15, 2024
Closed
weifengpy and others added 8 commits April 16, 2024 14:13
@weifengpy
Merge branch 'pytorch-labs:main' into main
3ac9d81
@weifengpy
UNIT TEST FOR STATE DICT
38461b3 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
@weifengpy
implement to
bc7a764 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
@weifengpy
remove torch.override from torch function
8b1d037 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
@weifengpy
use dtype in compile unit test
7ff6855 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
@weifengpy
add dtype in all unit test
d9bcf71 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
@weifengpy
keep original dtype
923bef2 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
@weifengpy
fix linter
e15d244 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
@weifengpy weifengpy changed the title [In Progress] FSDP2 + NF4Tensor [FSDP2][1/n] construct NF4Tensor from bf16/fp16/fp32 Apr 17, 2024
@weifengpy
use torch testing @parametrize
d4beb8f 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
@weifengpy
remove unused import
f41cb3d 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
@weifengpy
Copy link
Contributor Author

weifengpy commented Apr 17, 2024 •
edited
Loading

error test/kernel/test_galore_downproj.py seems to be irrelevant. it's throws error when parsing nvidia-smi. cc @msaroufim @jeromeku since it's newly added this week

@weifengpy weifengpy marked this pull request as ready for review April 17, 2024 22:03
@cpuhrsch
Copy link
Contributor

Great! Thank you! Can you try rebasing to see if CI runs green?

@weifengpy
Copy link
Contributor Author

Great! Thank you! Can you try rebasing to see if CI runs green?

rebasing now

drisspg
drisspg reviewed Apr 17, 2024
View reviewed changes
torchao/dtypes/nf4tensor.py Outdated

@implements_torch_function(torch.Tensor.to)
def function_to_dtype(*args, **kwargs):
return args[0].get_original_weight().to(args[1])
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit there is a few ways you can call to that this isn't robust too I would imagine

Copy link
Contributor Author

@weifengpy weifengpy Apr 17, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

do you mean 3 ways to call .to ? I can raise unimplemented if it helps

  • to(dtype, non_blocking=False, copy=False, memory_format=torch.preserve_format)
  • torch.to(device=None, dtype=None, non_blocking=False, copy=False, memory_format=torch.preserve_format)
  • torch.to(other, non_blocking=False, copy=False)

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

yeah, I think in theory they all can be supported its just that the other args/kwargs are getting dropped as implemented

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

got you. will see if I can pass in args/kwargs instead of dropping them

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

updated the PR to include args and kwargs instead of dropping. passing to dispatch if not implemented

drisspg
drisspg approved these changes Apr 17, 2024
View reviewed changes
Copy link
Contributor

@drisspg drisspg left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Overall looks good, thanks for pushing this through!

@weifengpy
sm8 for fp16
950d9fd 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
@weifengpy

This comment was marked as outdated.

Sign in to view
@weifengpy
remove sm check for fp16
d4eae0b 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
@weifengpy weifengpy marked this pull request as draft April 18, 2024 00:56
weifengpy and others added 4 commits April 18, 2024 16:12
@weifengpy
skip 2.2.2 and below for tracing tensor subclass
9444f2c 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
@weifengpy
Merge branch 'pytorch-labs:main' into main
b2c3c02
@weifengpy
include kwargs
9be2de3 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
@weifengpy
raise unimplemented
2981393 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
weifengpy
weifengpy commented Apr 19, 2024
View reviewed changes
test/dtypes/test_nf4.py
def test_smoketest_linear_compile(self):
for dtype in [torch.bfloat16, torch.float16]:
if torch.cuda.is_available() and torch.cuda.get_device_capability() < (8, 0) and dtype == torch.bfloat16:
self.skipTest("test requires SM capability of at least (8, 0).")
Copy link
Contributor Author

@weifengpy weifengpy Apr 19, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

test_smoketest_linear_compile is always skipped before because we exit with self.skipTest with torch.bfloat16 and did not have a chance to test torch.float16. It is fixed in this version by using @parameterize

weifengpy
weifengpy commented Apr 19, 2024
View reviewed changes
test/dtypes/test_nf4.py
if torch.cuda.is_available() and torch.cuda.get_device_capability() < (8, 0) and dtype == torch.bfloat16:
self.skipTest("test requires SM capability of at least (8, 0).")
if version.parse(torch.__version__) < version.parse("2.3.0"):
self.skipTest("test requires 2.3.0 and above for tracing NF4Tensor")
Copy link
Contributor Author

@weifengpy weifengpy Apr 19, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

starting from 2.3.0 we can trace subclass when inner tensors have different shapes than outer wrapper class. Specifically, we use symbolic_context.inner_contexts instead of symbolic_context from outer wrapper class: https:/pytorch/pytorch/blob/main/torch/_subclasses/meta_utils.py#L649

@weifengpy weifengpy marked this pull request as ready for review April 19, 2024 01:53
@cpuhrsch cpuhrsch merged commit a7ff835 into pytorch:main Apr 19, 2024
13 checks passed
@weifengpy weifengpy mentioned this pull request Apr 19, 2024
Closed
dbyoung18 pushed a commit to dbyoung18/ao that referenced this pull request Jul 31, 2024
@weifengpy
[FSDP2][1/n] construct NF4Tensor from bf16/fp16/fp32 (pytorch#118)
ac0dd2b
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@drisspg drisspg drisspg approved these changes

@cpuhrsch cpuhrsch cpuhrsch approved these changes

@rohan-varma rohan-varma Awaiting requested review from rohan-varma

Assignees
No one assigned
Labels
CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed.
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
4 participants
@weifengpy @cpuhrsch @drisspg @facebook-github-bot