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.

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

Add **quantization_kwargs to FrozenNF4Linear and LoRALinear and DoRALinear #1987

Merged
merged 4 commits into from
Nov 12, 2024

Conversation

joecummings
Copy link
Contributor

@joecummings joecummings commented Nov 12, 2024

Context

What is the purpose of this PR? Is it to

  • add a new feature
  • fix a bug
  • update tests and/or documentation
  • other (please add here)

The original error that sparked this investigation was that I was able to run Llama3.2V QLoRA on 4 GPUs, but it failed on 8 GPUs with the following error:

[rank0]:   File "/home/jrcummings/.conda/envs/tt-v0.4.0-rc2/lib/python3.11/site-packages/recipes/lora_finetune_distributed.py", line 274, in setup
[rank0]:     self._model = self._setup_model(
[rank0]:                   ^^^^^^^^^^^^^^^^^^
[rank0]:   File "/home/jrcummings/.conda/envs/tt-v0.4.0-rc2/lib/python3.11/site-packages/recipes/lora_finetune_distributed.py", line 473, in _setup_model
[rank0]:     training.shard_model(
[rank0]:   File "/home/jrcummings/.conda/envs/tt-v0.4.0-rc2/lib/python3.11/site-packages/torchtune/training/_distributed.py", line 643, in shard_model
[rank0]:     fully_shard(m, **fsdp_kwargs)
[rank0]:   File "/home/jrcummings/.conda/envs/tt-v0.4.0-rc2/lib/python3.11/site-packages/torch/distributed/_composable/contract.py", line 125, in wrapper
[rank0]:     updated = func(inp_module, *args, **kwargs)
[rank0]:               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/home/jrcummings/.conda/envs/tt-v0.4.0-rc2/lib/python3.11/site-packages/torch/distributed/_composable/fsdp/fully_shard.py", line 132, in fully_shard
[rank0]:     state._fsdp_param_group = FSDPParamGroup(
[rank0]:                               ^^^^^^^^^^^^^^^
[rank0]:   File "/home/jrcummings/.conda/envs/tt-v0.4.0-rc2/lib/python3.11/site-packages/torch/distributed/_composable/fsdp/_fsdp_param_group.py", line 114, in __init__
[rank0]:     self.fsdp_params = [
[rank0]:                        ^
[rank0]:   File "/home/jrcummings/.conda/envs/tt-v0.4.0-rc2/lib/python3.11/site-packages/torch/distributed/_composable/fsdp/_fsdp_param_group.py", line 115, in <listcomp>
[rank0]:     FSDPParam(
[rank0]:   File "/home/jrcummings/.conda/envs/tt-v0.4.0-rc2/lib/python3.11/site-packages/torch/distributed/_composable/fsdp/_fsdp_param.py", line 231, in __init__
[rank0]:     self._init_sharded_param(param, device)
[rank0]:   File "/home/jrcummings/.conda/envs/tt-v0.4.0-rc2/lib/python3.11/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
[rank0]:     return func(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/home/jrcummings/.conda/envs/tt-v0.4.0-rc2/lib/python3.11/site-packages/torch/distributed/_composable/fsdp/_fsdp_param.py", line 335, in _init_sharded_param
[rank0]:     chunks = _chunk_with_empty(param_data, shard_world_size, dim=0)
[rank0]:              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/home/jrcummings/.conda/envs/tt-v0.4.0-rc2/lib/python3.11/site-packages/torch/distributed/_composable/fsdp/_fsdp_common.py", line 95, in _chunk_with_empty
[rank0]:     chunks = list(torch.chunk(tensor, num_chunks, dim=dim))
[rank0]:                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/home/jrcummings/.conda/envs/tt-v0.4.0-rc2/lib/python3.11/site-packages/torchao/dtypes/nf4tensor.py", line 850, in __torch_function__
[rank0]:     return func(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/home/jrcummings/.conda/envs/tt-v0.4.0-rc2/lib/python3.11/site-packages/torch/_dynamo/eval_frame.py", line 632, in _fn
[rank0]:     return fn(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/home/jrcummings/.conda/envs/tt-v0.4.0-rc2/lib/python3.11/site-packages/torchao/dtypes/nf4tensor.py", line 831, in __torch_dispatch__
[rank0]:     return NF4_OPS_TABLE[func](func, args, kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/home/jrcummings/.conda/envs/tt-v0.4.0-rc2/lib/python3.11/site-packages/torchao/dtypes/nf4tensor.py", line 195, in nf4_split
[rank0]:     inner_tensor.numel() % num_chunks == 0
[rank0]: AssertionError: quantization_factor.numel() not divisible by 8

Weird, right? In digging further @pbontrager figured out that the default block_size and scaler_block_size used when converting the weights to NF4 were the wrong size when the model was sharded to 8 GPUs. After consulting with @drisspg, it seems relatively harmless to modify the scaler_block_size when quantizing these weights. Therefore, the fix here is to expose the quantization_kwargs for both FrozenNF4Linear and LoRALinear. A follow up PR will actually land the changes in the model builders to resolve the initial error.

Changelog

What are the changes made in this PR?

  • Add quantization_kwargs to FrozenNF4Linear
  • Add test for quantization kwargs and FrozenNF4Linear
  • Add quantization_kwargs to LoRALinear
  • Add tests for quantization kwargs and LoRALinear
  • Add quantization_kwargs to DoRALinear
  • Add tests for quantization kwargs and DoRALinear
  • Profit?

Test plan

Please make sure to do each of the following if applicable to your PR. If you're unsure about any one of these just ask and we will happily help. We also have a contributing page for some guidance on contributing.

  • run pre-commit hooks and linters (make sure you've first installed via pre-commit install)
  • add unit tests for any new functionality
  • update docstrings for any new or updated methods or classes
  • run unit tests via pytest tests
  • run recipe tests via pytest tests -m integration_test
  • manually run any new or modified recipes with sufficient proof of correctness
  • include relevant commands and any other artifacts in this summary (pastes of loss curves, eval results, etc.)

UX

If your function changed a public API, please add a dummy example of what the user experience will look like when calling it.
Here is a docstring example
and a tutorial example

  • I did not change any public API
  • I have added an example to docs or docstrings

Copy link

pytorch-bot bot commented Nov 12, 2024

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/torchtune/1987

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 1196f8a with merge base ac4f88e (image):
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@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 Nov 12, 2024
@codecov-commenter
Copy link

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 64.96%. Comparing base (ac4f88e) to head (e374fcd).

Additional details and impacted files
@@             Coverage Diff             @@
##             main    #1987       +/-   ##
===========================================
+ Coverage   24.79%   64.96%   +40.16%     
===========================================
  Files         318      318               
  Lines       17597    17610       +13     
===========================================
+ Hits         4364    11440     +7076     
+ Misses      13233     6170     -7063     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

@joecummings joecummings changed the title Add **quantization_kwargs to FrozenNF4Linear and LoRALinear Add **quantization_kwargs to FrozenNF4Linear and LoRALinear and DoRALinear Nov 12, 2024
@joecummings joecummings merged commit 4df97ad into pytorch:main Nov 12, 2024
17 checks passed
@joecummings joecummings deleted the update-quant-settings branch November 12, 2024 16:37
"""Test that passing in non-default quantization kwargs works as expected."""
quantization_kwargs = {
"block_size": 16,
"scaler_block_size": 256,
Copy link

Choose a reason for hiding this comment

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

Sad my spelling mistake is now surfaced lol

joecummings added a commit that referenced this pull request Nov 13, 2024
@ebsmothers ebsmothers mentioned this pull request Nov 26, 2024
44 tasks
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
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
Development

Successfully merging this pull request may close these issues.

5 participants