pytorch-labs / float8_experimental

This repository contains the experimental PyTorch native float8 training UX
BSD 3-Clause "New" or "Revised" License
194 stars 18 forks source link

[3/x]: simplify FSDP1 test and add coverage for dynamic scaling #293

Closed vkuzo closed 2 weeks ago

vkuzo commented 2 weeks ago

Stack from ghstack (oldest at bottom):

Summary:

  1. simplify the FSDP test, instead of testing 1 GPU vs N GPUs, instead hold the number of GPUs constant and test bf16 vs float8. Remove various technical debt that accumulated in this test.
  2. add testing for dynamic scaling of weights

Test Plan:

./test/test_fsdp.sh

Reviewers:

Subscribers:

Tasks:

Tags:

Differential Revision: D59305791

vkuzo commented 2 weeks ago

@vkuzo has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

facebook-github-bot commented 2 weeks ago

This pull request has been merged in pytorch-labs/float8_experimental@006f89aa4e6d057bdf9eabc1984448ad98edddfb.