Torchao float8 training #3348
Torchao float8 training #3348
Conversation
benchmarks/fp8/torchao/ddp.py
Outdated
| old_named_params = get_named_parameters(model) | ||
|
|
||
| with torch.no_grad(): | ||
| convert_model(model) |
There was a problem hiding this comment.
looks like this is still using TE - intended?
There was a problem hiding this comment.
Nope, needed to update these scripts, thought I had only pushed single-GPU! Though now DDP and FSDP are good, verifying DeepSpeed in a moment
| training with `torchao` FP8. | ||
|
|
||
| Args: | ||
| recipe_name (`str`, *optional*, default to `None`): |
There was a problem hiding this comment.
is it ok to leave the recipe by name lookup out for now? Technically this is a private API, we really like it internally and want to move it to public, but wanted to do a round of community feedback + design before we do that. The default recipe (Float8LinearConfig()) is a great starting point!
There was a problem hiding this comment.
Sure, will tweak accordingly
| @@ -142,6 +142,10 @@ def is_torch_xla_available(check_is_tpu=False, check_is_gpu=False): | |||
| return True | |||
|
|
|||
|
|
|||
| def is_torchao_available(): | |||
| return _is_package_available("torchao") | |||
There was a problem hiding this comment.
would it be possible to require torchao v0.6.1 or later? It should work with earlier versions but a more recent version is ideal for better performance. If needed I can look up the exact earliest version which we can expose - lmk.
There was a problem hiding this comment.
I see no issue with that, especially as it's a young library too which makes our lives easier since it's not in accelerate yet!
What does this PR do?
This PR is the start of incorporating torchao float8 training into Accelerate.
It does so by refactoring the backend to be more modular as more FP8 backends come out, and adds new benchmarking tests for it.
Current status of benchmarks:
Also need to work on a testing suite.
Fixes # (issue)
Before submitting
Pull Request section?
to it if that's the case.
documentation guidelines, and
here are tips on formatting docstrings.
Who can review?
Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.
@SunMarc