Fixes Issue #176 which caused torch.compile to graphbreak and fail with float8 (#177)

Summary:
# Summary

Attempts to fix #176

Current test_compile errors:
```Shell
                        (output, treespec),
                        graph,
                        lifted_freevars,
                    )

        except Unsupported as ex:
            f_name = f"{type(f).__name__}"
            if isinstance(f, UserFunctionVariable):
                f_name = f.get_name()
            msg = (
                f"speculate_subgraph: while introspecting {description}, we were unable "
                f"to trace function `{f_name}` into a single graph. This means "
                f"that Dynamo was unable to prove safety for this API and will "
                f"fall back to eager-mode PyTorch, which could lead to a slowdown."
            )
            log.warning(msg)
            log.exception(ex)
>           raise Unsupported(
                f"{msg} Scroll up for the stack trace "
                f"of the initial exception. The reason was: {ex.msg}"
            ) from ex
E           torch._dynamo.exc.Unsupported: speculate_subgraph: while introspecting autograd.Function, we were unable to trace function `backward` into a single graph. This means that Dynamo was unable to prove safety for this API and will fall back to eager-mode PyTorch, which could lead to a slowdown. Scroll up for the stack trace of the initial exception. The reason was: call_function UserDefinedClassVariable() [TensorVariable(), TensorVariable(), ConstantVariable(dtype)] {'emulate': ConstantVariable(bool)}
E
E           from user code:
E              File "/home/drisspg/meta/float8_experimental/float8_experimental/dynamic_linear/dynamic_float8_linear.py", line 62, in forward
E               y = self.cast_to_float8e5m2_bw(y)
E             File "/home/drisspg/meta/float8_experimental/float8_experimental/dynamic_linear/dynamic_float8_linear.py", line 79, in cast_to_float8e5m2_bw
E               return NoopFwToFloat8E5M2Bw.apply(gradY, self.emulate)
E
E           Set TORCH_LOGS="+dynamo" and TORCHDYNAMO_VERBOSE=1 for more information
E
E
E           You can suppress this exception and fall back to eager by setting:
E               import torch._dynamo
E               torch._dynamo.config.suppress_errors = True

../../miniconda3/envs/nightly/lib/python3.10/site-packages/torch/_dynamo/variables/higher_order_ops.py:328: Unsupported
---------------------------------------- Captured stderr call -----------------------------------------
[2024-01-02 15:43:30,115] [29/0] torch._dynamo.variables.higher_order_ops: [WARNING] speculate_subgraph: while introspecting autograd.Function, we were unable to trace function `backward` into a single graph. This means that Dynamo was unable to prove safety for this API and will fall back to eager-mode PyTorch, which could lead to a slowdown.
[2024-01-02 15:43:30,115] [29/0] torch._dynamo.variables.higher_order_ops: [ERROR] call_function UserDefinedClassVariable() [TensorVariable(), TensorVariable(), ConstantVariable(dtype)] {'emulate': ConstantVariable(bool)}
[2024-01-02 15:43:30,115] [29/0] torch._dynamo.variables.higher_order_ops: [ERROR] Traceback (most recent call last):
[2024-01-02 15:43:30,115] [29/0] torch._dynamo.variables.higher_order_ops: [ERROR]   File "/home/drisspg/miniconda3/envs/nightly/lib/python3.10/site-packages/torch/_dynamo/variables/higher_order_ops.py", line 256, in speculate_subgraph
[2024-01-02 15:43:30,115] [29/0] torch._dynamo.variables.higher_order_ops: [ERROR]     output = f.call_function(tx, args, sub_kwargs)
[2024-01-02 15:43:30,115] [29/0] torch._dynamo.variables.higher_order_ops: [ERROR]   File "/home/drisspg/miniconda3/envs/nightly/lib/python3.10/site-packages/torch/_dynamo/variables/functions.py", line 255, in call_function
[2024-01-02 15:43:30,115] [29/0] torch._dynamo.variables.higher_order_ops: [ERROR]     return super().call_function(tx, args, kwargs)
[2024-01-02 15:43:30,115] [29/0] torch._dynamo.variables.higher_order_ops: [ERROR]   File "/home/drisspg/miniconda3/envs/nightly/lib/python3.10/site-packages/torch/_dynamo/variables/functions.py", line 84, in call_function
[2024-01-02 15:43:30,115] [29/0] torch._dynamo.variables.higher_order_ops: [ERROR]     return tx.inline_user_function_return(
[2024-01-02 15:43:30,115] [29/0] torch._dynamo.variables.higher_order_ops: [ERROR]   File "/home/drisspg/miniconda3/envs/nightly/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 650, in inline_user_function_return
[2024-01-02 15:43:30,115] [29/0] torch._dynamo.variables.higher_order_ops: [ERROR]     return InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
[2024-01-02 15:43:30,115] [29/0] torch._dynamo.variables.higher_order_ops: [ERROR]   File "/home/drisspg/miniconda3/envs/nightly/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 2245, in inline_call
[2024-01-02 15:43:30,115] [29/0] torch._dynamo.variables.higher_order_ops: [ERROR]     return cls.inline_call_(parent, func, args, kwargs)
[2024-01-02 15:43:30,115] [29/0] torch._dynamo.variables.higher_order_ops: [ERROR]   File "/home/drisspg/miniconda3/envs/nightly/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 2354, in inline_call_

```

Pull Request resolved: #177

Reviewed By: vkuzo

Differential Revision: D52495888

Pulled By: drisspg

fbshipit-source-id: f7dd1ab8c4395f2e51043ef3526eae666299a500

Author: drisspg

float8_experimental/dynamic_linear/dynamic_float8_linear.py

@@ -13,6 +13,7 @@
 from float8_experimental.float8_utils import tensor_to_scale, to_fp8_saturated
 
 
+@torch._dynamo.allow_in_graph
 class NoopFwToFloat8E5M2Bw(torch.autograd.Function):
     """
     Forward: no-op

float8_experimental/float8_linear.py

@@ -61,6 +61,7 @@ def _maybe_initialize_amaxes_scales_for_float8_cast(
         scale.copy_(new_scale)
 
 
+@torch._dynamo.allow_in_graph
 class NoopFwToFloat8E5M2Bw(torch.autograd.Function):
     """
     Forward: no-op

float8_experimental/float8_tensor.py

@@ -22,6 +22,7 @@ def calculate_amax_and_cast_to_float8(tensor, scale, float8_dtype, amax_buffer):
     return bits_fp8
 
 
+@torch._dynamo.allow_in_graph
 class ToFloat8ConstrFunc(torch.autograd.Function):
     """
     A differentiable conversion to fp8
@@ -54,6 +55,7 @@ def backward(ctx, g):
             return g, None, None, None, None, None
 
 
+@torch._dynamo.allow_in_graph
 class FromFloat8ConstrFunc(torch.autograd.Function):
     """
     A differentiable conversion from fp8
@@ -189,8 +191,3 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None):
 
     # Do not force the Float8Tensor type on the returned tensor
     __torch_function__ = torch._C._disabled_torch_function_impl
-
-
-# In order for dynamo to successfuly trace our tensor subclass, we need
-# to be able to represent it in the graph.
-torch._dynamo.allow_in_graph(Float8Tensor)