Reduced CPU overhead in precompute_float8_dynamic_scale_for_fsdp

[awgu]

Stack from ghstack (oldest at bottom):

• -> Reduced CPU overhead in precompute_float8_dynamic_scale_for_fsdp #331

Description
For Llama3-8B on 8xH100 profiling with with_stack=True (which does add overhead), the precompute_float8_dynamic_scale_for_fsdp CPU time decreases from 24 ms to 15 ms.

Before:

After:

Test Plan

(pytorch-3.10) [[email protected] /data/users/andgu/float8_experimental (precompute_float8)]$ pytest test/test_fsdp2/test_fsdp2.py 
========================================================= test session starts =========================================================
platform linux -- Python 3.10.13, pytest-7.3.2, pluggy-1.3.0
rootdir: /data/users/andgu/float8_experimental
plugins: xdoctest-1.1.0, hypothesis-5.35.1, xdist-3.3.1, shard-0.1.2, rerunfailures-13.0, flakefinder-1.1.0, cpp-2.3.0
collected 8 items                                                                                                                     
Running 8 items in this shard

test/test_fsdp2/test_fsdp2.py ........                                                                                          [100%]

========================================================== warnings summary ===========================================================
test/test_fsdp2/test_fsdp2.py::TestFloat8MultiThread::test_fp32_fp8_multi_module_parity
test/test_fsdp2/test_fsdp2.py::TestFloat8MultiThread::test_fp32_fp8_single_module_parity
  /data/users/andgu/float8_experimental/float8_experimental/float8_linear_utils.py:272: FutureWarning: The combination of ranks + tag as process group identifier has been deprecated. Please switch to using ProcessGroup, DeviceMesh, or group name instead.
    all_reduced_amax_tensor = all_reduce(

-- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html
============================================== 8 passed, 2 warnings in 121.90s (0:02:01) ==============================================

Differential Revision: D60236258

[awgu]

We have to call to_local() here because DTensor does not support [i] int indexing. Int indexing might not be semantically clear if the DTensor is sharded; I think indexing should be on the local tensor.

[awgu]

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

[drisspg]

This is a TLDR for me, I like vstack because semantically I think about gluing lego blocks together lol Does it assert some contiguousnous that causes it to be less performant?

[awgu]

From offline: vstack incurs per-tensor reshape, which each redispatches through DTensor dispatch, which is what makes it slow. stack only goes through DTensor dispatch once.

[facebook-github-bot]

@awgu merged this pull request in 701647b.