[PyTorch] cuda graph support

[ksivaman]

This PR adds the following features (high-level):

• make_graphed_callables API similar to the PyTorch API with some additional arguments for FP8 usage. Support for fp8 weight caching via existing is_first_microbatchargument is also retained.
• Restructuring and amax reduction logic with a simpler design and handling of various parallelisms with minimal book-keeping compared to the previous approach.
• Forward and backward amaxes are reduced within the scope of current iteration, solving numerous bugs w.r.t. checkpointing and removing the need to save global buffers.
• Support for nested/multiple FP8 autocast contexts with different recipes and distributed groups.
• Amax reductions are module independent and happen at at autocast level. This also resolves numerous bugs and allows for support for MoE/LoRA like models.
• Redesign of transposes for Float8Tensor that makes the transposes persistent for graph capture. Also fixes use cases for the vanilla optimizers (non fp8-distopt).
• The scaling inverses for weight tensors are no longer frozen when caching weights across microbatches.

[timmoon10]

Removing the automatic cache clearing makes using the transpose cache a much more manual and dangerous process. Consider something like:

matmul(x, w.transpose(0, 1))
w -= learning_rate * w.grad
matmul(x, w.transpose(0, 1))

Previously we could just set update_cache="lazy". Now there needs to be manual logic to figure out the microbatch step, or else it will provide the stale values.

[ksivaman]

In this example, caching is not used, so a fresh transpose will be computed each time.

[ksivaman]

If caching is used, it is reasonable to expect the user to know when to reuse a cached value and when to force recompute. This is consistent with our design of is_first_microbatch argument to the forward for module APIs.

[ksivaman]

Note: we use 2 args cache and update_cache to support this logic.

[timmoon10]

I think we're overfitting to the Linear weight use-case. For example, in #707 I want to pass Float8Tensors between ops as inputs or dgrads:

class DbiasCastTranspose:
    def backward(self, dy):
        db = dy.sum(dim=0)
        dx = cast_transpose(dy)  # Creates Float8Tensor with transpose cache
        return dx, db

class FP8Linear:  # Part of FP8 attention
    def backward(self, dy):
        if not isinstance(dy, Float8Tensor):
           dy = Float8Tensor.to_float8(dy)
        dx = Float8Tensor(...)  # No transpose cache
        fp8_gemm(w.transpose()._data, dy.transpose()._data, out=dx._data)
        dw = fp8_gemm(x, dy)
        return dx, dw

FP8Linear has no idea where its input came from. Maybe it's from DbiasCastTranspose (Float8Tensor with cached transpose), FP8Linear (Float8Tensor without cached transpose), or a non-FP8 op. Our current approach with lazy transpose caching gives us a lot of flexibility and I think we should abandon it only when really necessary.

I suppose this is not precisely relevant since it doesn't involve in-place operations, but a more general statement about the design of Float8Tensor.

[ksivaman]

/te-ci

[ksivaman]

/te-ci pytorch

[ksivaman]

/te-ci pytorch

[timmoon10]

#735 has some improvements to the Float8Tensor transpose function, which should reduce the divergence with #707. If there are no issues, we should merge that branch into this PR.

[ksivaman]

/te-ci pytorch

[ksivaman]

/te-ci pytorch

[ksivaman]

/te-ci pytorch

[ksivaman]

/te-ci pytorch

[ksivaman]

/te-ci pytorch