Updates to sync_float_amax_history

[drisspg]

Summary

Update docs, make sure this is friendly to dynamo

Perf

PyTorch Version	float8 Version	Eager Iterations per Second	Compile
Nightly	Main	1.15 it/s	2.10 it/s
Nightly	This PR	1.16 it/s	2.27 it/s

Trace	Compile URL	Eager
This PR	https://fburl.com/753ztao4	https://fburl.com/34yftzao
Main	https://fburl.com/a0gh9iof	https://fburl.com/u9c4ilmp

Things I have done/changed

Commit 1

• We previously had an fp8_classes argument that would be passed in, this was to enable working with the separate TP/SP classes, since we plan to have Dtensor be the solution I am removing for now.
• I put the child.amax_and_scale_synced module mutation under the enable_amax_init flag, this seemed to be causing graphbreaks cause of the module mutation

Commit 2

• We previously had all the history buffers be scaler tensors. This meant that to construct the combined tensor we needed to call torch.Tensor which was causing a HtoD sync under torch.compile. I needed to added a single dimension of size 1 and pipe that through all the places.
• Note that this meant we needed to update the to_hp to send back to original precision because line the scale upcasts the _data tensor

Commit 3

• Rewrote the sync function to do the torch.roll() on all the histories at once - side note not sure if this is more expensive than to clones since we really dont care about the wrapping behavior
• Same for generating the new scales from the grouped histories

Things to do

• There is still two loops and those are for mutating the the actual module values, not sure if there is another way around this..
• Going to try the functional collectives

[awgu]

I left some initial comments just scanning through the code.

[awgu]

Aside: It seems like a bug on AsyncCollectiveTensor if it does not make sure to call wait() when you torch.split it. 🤔

[bdhirsh]

Hmm...

I actually tried to add an optimization to AsyncCollectiveTensor a while ago: pytorch/pytorch#105240. Where the thought behind it was:

(1) if we have a collective in-flight, we only need to sync the collective if the data needs to be used in an operation

(2) view ops generally don't need to access the data of their inputs, only the metadata.

So that PR "delays" syncs when there are view ops called on AsyncCollectiveTensor (e.g. aten.split()).

[drisspg]

Indeed, I was working through this with @wanchaol and without the all_reduced_amax_tensor was garbage...

[bdhirsh]

I'm curious why the manual sync here is needed - does it paper over some underlying bug? Didn't see your last comment. I wonder if this is related to some internal AsyncCollectiveTensor issues 🤔

Let me know if either of you want to stare at it together (I think fixing the underlying bug will help other areas too)

[awgu]

Interestingly, I was not able to repro the issue: P1184812077

_foreach_copy_ and functional all-reduce without .wait() seemed to be fine, but I might not be repro-ing correctly.

[drisspg]

hmmm let me try again then

[drisspg]

You are right... however:

    return self.__get_result()
  File "/home/drisspg/miniconda3/envs/dev/lib/python3.10/concurrent/futures/_base.py", line 403, in __get_result
    raise self._exception
torch._dynamo.exc.BackendCompilerFailed: backend='inductor' raised:
CompilationError: at 10:30:def triton_(in_ptr0, in_ptr1, in_ptr2, in_ptr3, out_ptr0, out_ptr1, out_ptr3, out_ptr4, out_ptr5, out_ptr7, out_ptr8, out_ptr9, out_ptr11):
    xpid = tl.program_id(0)
    XBLOCK: tl.constexpr = 1024
    if xpid >= 0 and xpid < 1:
        xpid_offset = xpid - 0
        xnumel = 1
        xoffset = xpid_offset * XBLOCK
        xindex = xoffset + tl.arange(0, XBLOCK)[:, None]
        xmask = xindex < xnumel
        rindex = tl.arange(0, RBLOCK)[None, :]
                              ^
NameError('RBLOCK is not defined')

Set TORCH_LOGS="+dynamo" and TORCHDYNAMO_VERBOSE=1 for more information


You can suppress this exception and fall back to eager by setting:
    import torch._dynamo
    torch._dynamo.config.suppress_errors = True

compiling the sync function now causes an inductor fault... so I will create another PR and investigate there but otherwise I think thats everything for this one

[drisspg]

#213

[facebook-github-bot]

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

[awgu]

LGTM!

[awgu]

For my understanding, how does passing float8_layers affect torch.compile?

[drisspg]

honestly I wrote this at the beginning of the PR and now I am not sure that this is true, but I found the get_attr was causing graph breaks but that was a few iterations ago and reading the graph_breaks logs can be a lil hard to decipher some times

[facebook-github-bot]

@drisspg merged this pull request in 956195b.