FLUX FP8 Training Error

[faruknane]

Hi,

I can't train the FLUX model in float8 using torch ao. The sanity check (inference) is done successfully. That means I am able to do 25 step inference in float8 and generate a good quality image. However, when it comes to the training / the backward process, the code gives me error. I can't fix it. Could you please identify the issue here?

A small code script regarding how the model is defined:

def LoadPipeline(dtype = torch.bfloat16):
    
    pipeline = FluxKontextPipeline.from_pretrained("/home/cropy/flux_kontext", 
                                            local_files_only=True,
                                            # quantization_config=pipeline_quant_config,
                                            torch_dtype=torch.bfloat16)

    pipeline.to("cuda")


    # quantize_(
    #     pipeline.transformer,
    #     float8_dynamic_activation_float8_weight(),
    # )

    quantize_(
        pipeline.vae,
        float8_dynamic_activation_float8_weight(),
    )
    quantize_(
        pipeline.text_encoder,
        float8_dynamic_activation_float8_weight(),
    )
    quantize_(
        pipeline.text_encoder_2,
        float8_dynamic_activation_float8_weight(),
    )
    
    pipeline.vae = torch.compile(pipeline.vae, mode="max-autotune", fullgraph=True)
    pipeline.text_encoder = torch.compile(pipeline.text_encoder, mode="max-autotune", fullgraph=True)
    pipeline.text_encoder_2 = torch.compile(pipeline.text_encoder_2, mode="max-autotune", fullgraph=True)


    return pipeline

----- My Model -----
self.pipeline = LoadPipeline(self.target_dtype)
self.model = self.pipeline.transformer.to(self.target_dtype)


train_fp8_config = Float8LinearConfig.from_recipe_name("rowwise") # tried tensorwise as well
convert_to_float8_training(self.model, config=train_fp8_config, module_filter_fn=module_filter_fn)

The error:

E1002 02:15:43.590000 278392 site-packages/torch/_subclasses/fake_tensor.py:2721] [2/1] failed while attempting to run meta for aten._scaled_mm.default
E1002 02:15:43.590000 278392 site-packages/torch/_subclasses/fake_tensor.py:2721] [2/1] Traceback (most recent call last):
E1002 02:15:43.590000 278392 site-packages/torch/_subclasses/fake_tensor.py:2721] [2/1]   File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torch/_subclasses/fake_tensor.py", line 2717, in _dispatch_impl
E1002 02:15:43.590000 278392 site-packages/torch/_subclasses/fake_tensor.py:2721] [2/1]     r = func(*args, **kwargs)
E1002 02:15:43.590000 278392 site-packages/torch/_subclasses/fake_tensor.py:2721] [2/1]         ^^^^^^^^^^^^^^^^^^^^^
E1002 02:15:43.590000 278392 site-packages/torch/_subclasses/fake_tensor.py:2721] [2/1]   File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torch/_ops.py", line 829, in __call__
E1002 02:15:43.590000 278392 site-packages/torch/_subclasses/fake_tensor.py:2721] [2/1]     return self._op(*args, **kwargs)
E1002 02:15:43.590000 278392 site-packages/torch/_subclasses/fake_tensor.py:2721] [2/1]            ^^^^^^^^^^^^^^^^^^^^^^^^^
E1002 02:15:43.590000 278392 site-packages/torch/_subclasses/fake_tensor.py:2721] [2/1]   File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torch/_meta_registrations.py", line 6448, in meta_scaled_mm
E1002 02:15:43.590000 278392 site-packages/torch/_subclasses/fake_tensor.py:2721] [2/1]     torch._check(
E1002 02:15:43.590000 278392 site-packages/torch/_subclasses/fake_tensor.py:2721] [2/1]   File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torch/__init__.py", line 1684, in _check
E1002 02:15:43.590000 278392 site-packages/torch/_subclasses/fake_tensor.py:2721] [2/1]     _check_with(RuntimeError, cond, message)
E1002 02:15:43.590000 278392 site-packages/torch/_subclasses/fake_tensor.py:2721] [2/1]   File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torch/__init__.py", line 1666, in _check_with
E1002 02:15:43.590000 278392 site-packages/torch/_subclasses/fake_tensor.py:2721] [2/1]     raise error_type(message_evaluated)
E1002 02:15:43.590000 278392 site-packages/torch/_subclasses/fake_tensor.py:2721] [2/1] RuntimeError: self must be row_major, got stride (1, 18432)
/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torch/autograd/graph.py:829: UserWarning: Error detected in matmul_with_hp_or_float8_argsBackward. Traceback of forward call that caused the error:
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/diffusers/models/transformers/transformer_flux.py", line 537, in forward
    return torch.utils.checkpoint.checkpoint(
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/diffusers/models/transformers/transformer_flux.py", line 481, in _forward
    norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.norm1_context(
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/diffusers/models/normalization.py", line 168, in forward
    emb = self.linear(self.silu(emb))
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torchao/float8/float8_linear.py", line 264, in forward
    output = matmul_with_hp_or_float8_args.apply(
 (Triggered internally at /pytorch/torch/csrc/autograd/python_anomaly_mode.cpp:122.)
  return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
Traceback (most recent call last):
  File "/mnt/c/Users/user_1/Desktop/rectifiedflow_main_image_editing/rectifiedflow_main_image_editing_flux/main.py", line 155, in <module>
    trainer.fit(model, train_dataloader, val_dataloader)
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/lightning/pytorch/trainer/trainer.py", line 561, in fit
    call._call_and_handle_interrupt(
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/lightning/pytorch/trainer/call.py", line 48, in _call_and_handle_interrupt
    return trainer_fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/lightning/pytorch/trainer/trainer.py", line 599, in _fit_impl
    self._run(model, ckpt_path=ckpt_path)
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/lightning/pytorch/trainer/trainer.py", line 1012, in _run
    results = self._run_stage()
              ^^^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/lightning/pytorch/trainer/trainer.py", line 1056, in _run_stage
    self.fit_loop.run()
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/lightning/pytorch/loops/fit_loop.py", line 216, in run
    self.advance()
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/lightning/pytorch/loops/fit_loop.py", line 455, in advance
    self.epoch_loop.run(self._data_fetcher)
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/lightning/pytorch/loops/training_epoch_loop.py", line 150, in run
    self.advance(data_fetcher)
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/lightning/pytorch/loops/training_epoch_loop.py", line 320, in advance
    batch_output = self.automatic_optimization.run(trainer.optimizers[0], batch_idx, kwargs)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/lightning/pytorch/loops/optimization/automatic.py", line 192, in run
    self._optimizer_step(batch_idx, closure)
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/lightning/pytorch/loops/optimization/automatic.py", line 270, in _optimizer_step
    call._call_lightning_module_hook(
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/lightning/pytorch/trainer/call.py", line 176, in _call_lightning_module_hook
    output = fn(*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/lightning/pytorch/core/module.py", line 1302, in optimizer_step
    optimizer.step(closure=optimizer_closure)
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/lightning/pytorch/core/optimizer.py", line 154, in step
    step_output = self._strategy.optimizer_step(self._optimizer, closure, **kwargs)
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/lightning/pytorch/strategies/strategy.py", line 239, in optimizer_step
    return self.precision_plugin.optimizer_step(optimizer, model=model, closure=closure, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/lightning/pytorch/plugins/precision/amp.py", line 76, in optimizer_step
    return super().optimizer_step(optimizer, model=model, closure=closure, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/lightning/pytorch/plugins/precision/precision.py", line 123, in optimizer_step
    return optimizer.step(closure=closure, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torch/optim/lr_scheduler.py", line 133, in wrapper
    return func.__get__(opt, opt.__class__)(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torch/optim/optimizer.py", line 516, in wrapper
    out = func(*args, **kwargs)
          ^^^^^^^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torch/utils/_contextlib.py", line 120, in decorate_context
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/bitsandbytes/optim/optimizer.py", line 272, in step
    loss = closure()
           ^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/lightning/pytorch/plugins/precision/precision.py", line 109, in _wrap_closure
    closure_result = closure()
                     ^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/lightning/pytorch/loops/optimization/automatic.py", line 146, in __call__
    self._result = self.closure(*args, **kwargs)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torch/utils/_contextlib.py", line 120, in decorate_context
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/lightning/pytorch/loops/optimization/automatic.py", line 131, in closure
    step_output = self._step_fn()
                  ^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/lightning/pytorch/loops/optimization/automatic.py", line 319, in _training_step
    training_step_output = call._call_strategy_hook(trainer, "training_step", *kwargs.values())
                           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/lightning/pytorch/trainer/call.py", line 328, in _call_strategy_hook
    output = fn(*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/lightning/pytorch/strategies/strategy.py", line 391, in training_step
    return self.lightning_module.training_step(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/mnt/c/Users/user_1/Desktop/rectifiedflow_main_image_editing/rectifiedflow_main_image_editing_flux/models/models.py", line 71, in training_step
    loss = self.flow.loss_fn(self.apply_model, **params)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/mnt/c/Users/user_1/Desktop/rectifiedflow_main_image_editing/rectifiedflow_main_image_editing_flux/modules/diffusion/flows.py", line 46, in loss_fn
    pred = model_fn(xt, t=t, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/mnt/c/Users/user_1/Desktop/rectifiedflow_main_image_editing/rectifiedflow_main_image_editing_flux/models/models.py", line 440, in apply_model
    v_pred = m(hidden_states=hidden_states, timestep=t, return_dict=False, **kwargs2)[0].clone()
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1784, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/diffusers/models/transformers/transformer_flux.py", line 802, in forward
    encoder_hidden_states, hidden_states = block(
                                           ^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1784, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torch/_dynamo/eval_frame.py", line 749, in compile_wrapper
    raise e.remove_dynamo_frames() from None  # see TORCHDYNAMO_VERBOSE=1
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torch/_dynamo/output_graph.py", line 1871, in _call_user_compiler
    raise BackendCompilerFailed(
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torch/_dynamo/output_graph.py", line 1846, in _call_user_compiler
    compiled_fn = compiler_fn(gm, example_inputs)
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torch/_dynamo/repro/after_dynamo.py", line 150, in __call__
    compiled_gm = compiler_fn(gm, example_inputs)
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torch/__init__.py", line 2380, in __call__
    return compile_fx(model_, inputs_, config_patches=self.config)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torch/_inductor/compile_fx.py", line 2418, in compile_fx
    return aot_autograd(
           ^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torch/_dynamo/backends/common.py", line 109, in __call__
    cg = aot_module_simplified(gm, example_inputs, **self.kwargs)
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torch/_functorch/aot_autograd.py", line 1199, in aot_module_simplified
    compiled_fn = AOTAutogradCache.load(
                  ^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torch/_functorch/_aot_autograd/autograd_cache.py", line 1140, in load
    compiled_fn = dispatch_and_compile()
                  ^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torch/_functorch/aot_autograd.py", line 1184, in dispatch_and_compile
    compiled_fn, _ = create_aot_dispatcher_function(
                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torch/_functorch/aot_autograd.py", line 576, in create_aot_dispatcher_function
    return _create_aot_dispatcher_function(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torch/_functorch/aot_autograd.py", line 836, in _create_aot_dispatcher_function
    compiled_fn, fw_metadata = compiler_fn(
                               ^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py", line 1262, in aot_dispatch_autograd
    fx_g, joint_inputs, maybe_subclass_meta = aot_dispatch_autograd_graph(
                                              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torch/_functorch/_aot_autograd/dispatch_and_compile_graph.py", line 318, in aot_dispatch_autograd_graph
    fx_g = _create_graph(joint_fn_to_trace, updated_joint_inputs, aot_config=aot_config)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torch/_functorch/_aot_autograd/dispatch_and_compile_graph.py", line 55, in _create_graph
    fx_g = make_fx(
           ^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py", line 2318, in wrapped
    return make_fx_tracer.trace(f, *args)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py", line 2250, in trace
    return self._trace_inner(f, *args)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py", line 2221, in _trace_inner
    t = dispatch_trace(
        ^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torch/_compile.py", line 53, in inner
    return disable_fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torch/_dynamo/eval_frame.py", line 929, in _fn
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py", line 1254, in dispatch_trace
    graph = tracer.trace(root, concrete_args)  # type: ignore[arg-type]
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torch/_dynamo/eval_frame.py", line 929, in _fn
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torch/fx/_symbolic_trace.py", line 850, in trace
    (self.create_arg(fn(*args)),),
                     ^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torch/fx/_symbolic_trace.py", line 703, in flatten_fn
    tree_out = root_fn(*tree_args)
               ^^^^^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py", line 1312, in wrapped
    out = f(*tensors)  # type:ignore[call-arg]
          ^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torch/_functorch/_aot_autograd/traced_function_transforms.py", line 720, in inner_fn
    outs = fn(*args)
           ^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torch/_functorch/_aot_autograd/traced_function_transforms.py", line 671, in joint_helper
    return _functionalized_f_helper(primals, tangents)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torch/_functorch/_aot_autograd/traced_function_transforms.py", line 419, in _functionalized_f_helper
    f_outs = fn(*f_args)
             ^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torch/_functorch/_aot_autograd/traced_function_transforms.py", line 286, in inner_fn_with_anomaly
    return inner_fn(*args)
           ^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torch/_functorch/_aot_autograd/traced_function_transforms.py", line 271, in inner_fn
    backward_out = torch.autograd.grad(
                   ^^^^^^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torch/autograd/__init__.py", line 452, in grad
    return handle_torch_function(
           ^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torch/overrides.py", line 1725, in handle_torch_function
    result = mode.__torch_function__(public_api, types, args, kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py", line 1360, in __torch_function__
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torch/autograd/__init__.py", line 503, in grad
    result = _engine_run_backward(
             ^^^^^^^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torch/autograd/graph.py", line 829, in _engine_run_backward
    return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torch/autograd/function.py", line 311, in apply
    return user_fn(self, *args)
           ^^^^^^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torchao/float8/float8_linear.py", line 197, in backward
    grad_weight = torch.mm(
                  ^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torchao/float8/float8_training_tensor.py", line 374, in __torch_dispatch__
    return FLOAT8_OPS_TABLE[func](func, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torchao/float8/float8_ops.py", line 385, in float8_mm
    tensor_out = addmm_float8_unwrapped(
                 ^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torchao/float8/float8_ops.py", line 72, in addmm_float8_unwrapped
    output = torch._scaled_mm(
             ^^^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torch/_subclasses/functional_tensor.py", line 511, in __torch_dispatch__
    outs_unwrapped = func._op_dk(
                     ^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torch/utils/_stats.py", line 28, in wrapper
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py", line 1462, in __torch_dispatch__
    return proxy_call(self, func, self.pre_dispatch, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py", line 914, in proxy_call
    out = func(*args, **kwargs)
          ^^^^^^^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torch/_ops.py", line 829, in __call__
    return self._op(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torch/utils/_stats.py", line 28, in wrapper
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torch/_subclasses/fake_tensor.py", line 1352, in __torch_dispatch__
    return self.dispatch(func, types, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torch/_subclasses/fake_tensor.py", line 2058, in dispatch
    return self._cached_dispatch_impl(func, types, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torch/_subclasses/fake_tensor.py", line 1487, in _cached_dispatch_impl
    output = self._dispatch_impl(func, types, args, kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torch/_subclasses/fake_tensor.py", line 2717, in _dispatch_impl
    r = func(*args, **kwargs)
        ^^^^^^^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torch/_ops.py", line 829, in __call__
    return self._op(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torch/_meta_registrations.py", line 6448, in meta_scaled_mm
    torch._check(
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torch/__init__.py", line 1684, in _check
    _check_with(RuntimeError, cond, message)
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torch/__init__.py", line 1666, in _check_with
    raise error_type(message_evaluated)
torch._dynamo.exc.BackendCompilerFailed: backend='inductor' raised:
RuntimeError: self must be row_major, got stride (1, 18432)

Set TORCHDYNAMO_VERBOSE=1 for the internal stack trace (please do this especially if you're reporting a bug to PyTorch). For even more developer context, set TORCH_LOGS="+dynamo"

I tried to make every input tensor contiguous for nn.Linear layers using the code below, but it didn't solve the issue.

class SafeFloat8Linear(nn.Module):
    def __init__(self, float8_linear):
        super().__init__()
        self.inner = float8_linear

    def forward(self, x, *args, **kwargs):
        return self.inner(x.contiguous(), *args, **kwargs)


def wrap_float8_linears(module: nn.Module):
    """
    Recursively replace all nn.Linear modules with SafeFloat8Linear-wrapped versions.
    """
    for name, child in list(module.named_children()):
        # If it's a Linear (or a specific float8 Linear class), wrap it
        if isinstance(child, nn.Linear):
            setattr(module, name, SafeFloat8Linear(child))
        else:
            # Recurse into children
            wrap_float8_linears(child)
    return module

[danielvegamyhre]

1. Which version of pytorch are you using?
2. Are you using torch.compile?

There was an issue with torch.compile ignoring strides (mem layout conversion), there was a fix to get compile and eager strides to match (can't find it off hand). I would retry this with torch nightly build.

[faruknane]

Hi @danielvegamyhre ,

My torch version: 2.8.0+cu128
I can't use direct torch compile on the model since it gives me an error (not related to fp8) about tensor being overwritten that I need to clone it or mark cuda step begin etc. but when I do those the issue won't go. So instead of direct model compilation, I compile the forward method of the inner flux attention layers like this:

@maybe_allow_in_graph
class FluxSingleTransformerBlock(nn.Module):
    def __init__(self, dim: int, num_attention_heads: int, attention_head_dim: int, mlp_ratio: float = 4.0, gradient_checkpointing = False):
        super().__init__()
        self.mlp_hidden_dim = int(dim * mlp_ratio)

        self.norm = AdaLayerNormZeroSingle(dim)
        self.proj_mlp = nn.Linear(dim, self.mlp_hidden_dim)
        self.act_mlp = nn.GELU(approximate="tanh")
        self.proj_out = nn.Linear(dim + self.mlp_hidden_dim, dim)

        self.attn = FluxAttention(
            query_dim=dim,
            dim_head=attention_head_dim,
            heads=num_attention_heads,
            out_dim=dim,
            bias=True,
            processor=FluxAttnProcessor(),
            eps=1e-6,
            pre_only=True,
        )

        self.gradient_checkpointing = gradient_checkpointing


    
    def _forward(
        self,
        hidden_states: torch.Tensor,
        encoder_hidden_states: torch.Tensor,
        temb: torch.Tensor,
        image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
        joint_attention_kwargs: Optional[Dict[str, Any]] = None,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        text_seq_len = encoder_hidden_states.shape[1]
        hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)

        residual = hidden_states
        norm_hidden_states, gate = self.norm(hidden_states, emb=temb)
        mlp_hidden_states = self.act_mlp(self.proj_mlp(norm_hidden_states))
        joint_attention_kwargs = joint_attention_kwargs or {}
        attn_output = self.attn(
            hidden_states=norm_hidden_states,
            image_rotary_emb=image_rotary_emb,
            **joint_attention_kwargs,
        )

        hidden_states = torch.cat([attn_output, mlp_hidden_states], dim=2)
        gate = gate.unsqueeze(1)
        hidden_states = gate * self.proj_out(hidden_states)
        hidden_states = residual + hidden_states
        if hidden_states.dtype == torch.float16:
            hidden_states = hidden_states.clip(-65504, 65504)

        encoder_hidden_states, hidden_states = hidden_states[:, :text_seq_len], hidden_states[:, text_seq_len:]
        return encoder_hidden_states, hidden_states

    @torch.compile
    def forward(
        self,
        hidden_states: torch.Tensor,
        encoder_hidden_states: torch.Tensor,
        temb: torch.Tensor,
        image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
        joint_attention_kwargs: Optional[Dict[str, Any]] = None,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        if self.gradient_checkpointing:
            return torch.utils.checkpoint.checkpoint(
                self._forward,
                hidden_states,
                encoder_hidden_states,
                temb,
                image_rotary_emb,
                joint_attention_kwargs,
            )
        else:
            return self._forward(
                hidden_states,
                encoder_hidden_states,
                temb,
                image_rotary_emb,
                joint_attention_kwargs,
            )

and I'm enabling gradient checkpointing if it does matter.

self.single_transformer_blocks = nn.ModuleList(
            [
                FluxSingleTransformerBlock(
                    dim=self.inner_dim,
                    num_attention_heads=num_attention_heads,
                    attention_head_dim=attention_head_dim,
                    gradient_checkpointing = True
                )
                for _ in range(num_single_layers)
            ]
        )

[danielvegamyhre]

If you could provide a self-contained repro script, that would be helpful.

Also, please be aware float8 training requires torch.compile for high performance, otherwise, performance may be worse than high precision baseline.

[faruknane]

@danielvegamyhre Have tried using torch 2.10.0 dev version. Uninstalled flash-attn which was giving me an error with torch 2.10.0. I am still getting the same error as before.

Torch Version:
Name: torch
Version: 2.10.0.dev20250930+cu128
Summary: Tensors and Dynamic neural networks in Python with strong GPU acceleration
Home-page: https://pytorch.org
Author:
Author-email: PyTorch Team packages@pytorch.org
License: BSD-3-Clause

Here is 3am. I will try to provide a sample code tomorrow for a complete picture.

[faruknane]

Hey @danielvegamyhre ,

Please check https://github.com/faruknane/TorchAO_BugReproduction for the bug reproduction. I added some notes to the readme file. Let me know if you have any questions or are able to produce the bug.

[danielvegamyhre]

Hey @danielvegamyhre ,

Please check https://github.com/faruknane/TorchAO_BugReproduction for the bug reproduction. I added some notes to the readme file. Let me know if you have any questions or are able to produce the bug.

Thanks for the repro example! In the notes you mention there are 3 changes need to be done to the code before running - could you make those changes in the code, so I can be 100% sure I'm running the same code as you?

[faruknane]

@danielvegamyhre The edits were made in the code of FLUX model which is in the transformers or diffusers library originally. Since, (all classes etc) they are all connected to each other, I m not able to do the changes on my code. I can share the edited file with you and maybe you could change the file in the library. Would that be ok?

[danielvegamyhre]

• Does this issue repro when you don't edit the transformers library code?
• Does this issue repro in eager mode?

[faruknane]

@danielvegamyhre I will try it without torch compile in 30 minutes. I didn't try without compile since I thought the issue was with torch compile. I'll let you know soon.

There was an issue with torch.compile ignoring strides (mem layout conversion), there was a fix to get compile and eager strides to match (can't find it off hand). I would retry this with torch nightly build

[faruknane]

@danielvegamyhre
I turned torch compile off. I still got an error but it was slightly different:

• without torch compile
• grad checkpointing = True

Sanity Checking DataLoader 0:   0%|                                                                                                                                                                                       | 0/2 [00:00<?, ?it/s]/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torch/_dynamo/eval_frame.py:929: UserWarning: torch.utils.checkpoint: the use_reentrant parameter should be passed explicitly. In version 2.5 we will raise an exception if use_reentrant is not passed. use_reentrant=False is recommended, but if you need to preserve the current default behavior, you can pass use_reentrant=True. Refer to docs for more details on the differences between the two variants.
  return fn(*args, **kwargs)
/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torch/utils/checkpoint.py:85: UserWarning: None of the inputs have requires_grad=True. Gradients will be None
  warnings.warn(
Sanity Checking DataLoader 0:  50%|███████████████████████████████████████████████████████████████████████████████████████▌                                                                                       | 1/2 [00:09<00:09,  0.10it/s]/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torch/_dynamo/eval_frame.py:929: UserWarning: torch.utils.checkpoint: the use_reentrant parameter should be passed explicitly. In version 2.5 we will raise an exception if use_reentrant is not passed. use_reentrant=False is recommended, but if you need to preserve the current default behavior, you can pass use_reentrant=True. Refer to docs for more details on the differences between the two variants.
  return fn(*args, **kwargs)
/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torch/utils/checkpoint.py:85: UserWarning: None of the inputs have requires_grad=True. Gradients will be None
  warnings.warn(
huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...                                                                                               
To disable this warning, you can either:
        - Avoid using `tokenizers` before the fork if possible
        - Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
To disable this warning, you can either:
        - Avoid using `tokenizers` before the fork if possible
        - Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
Epoch 0:   0%|                                                                                                                                                                                                         | 0/1000 [00:00<?, ?it/s]Epoch: 0 Training mode: True
Traceback (most recent call last):
  File "/mnt/c/Users/user_1/Desktop/rectifiedflow_main_image_editing/reproduce_error/main.py", line 68, in <module>
    trainer.fit(model, train_dataloader, train_dataloader)
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/lightning/pytorch/trainer/trainer.py", line 561, in fit
    call._call_and_handle_interrupt(
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/lightning/pytorch/trainer/call.py", line 48, in _call_and_handle_interrupt
    return trainer_fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/lightning/pytorch/trainer/trainer.py", line 599, in _fit_impl
    self._run(model, ckpt_path=ckpt_path)
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/lightning/pytorch/trainer/trainer.py", line 1012, in _run
    results = self._run_stage()
              ^^^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/lightning/pytorch/trainer/trainer.py", line 1056, in _run_stage
    self.fit_loop.run()
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/lightning/pytorch/loops/fit_loop.py", line 216, in run
    self.advance()
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/lightning/pytorch/loops/fit_loop.py", line 455, in advance
    self.epoch_loop.run(self._data_fetcher)
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/lightning/pytorch/loops/training_epoch_loop.py", line 150, in run
    self.advance(data_fetcher)
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/lightning/pytorch/loops/training_epoch_loop.py", line 320, in advance
    batch_output = self.automatic_optimization.run(trainer.optimizers[0], batch_idx, kwargs)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/lightning/pytorch/loops/optimization/automatic.py", line 192, in run
    self._optimizer_step(batch_idx, closure)
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/lightning/pytorch/loops/optimization/automatic.py", line 270, in _optimizer_step
    call._call_lightning_module_hook(
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/lightning/pytorch/trainer/call.py", line 176, in _call_lightning_module_hook
    output = fn(*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/lightning/pytorch/core/module.py", line 1302, in optimizer_step
    optimizer.step(closure=optimizer_closure)
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/lightning/pytorch/core/optimizer.py", line 154, in step
    step_output = self._strategy.optimizer_step(self._optimizer, closure, **kwargs)
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/lightning/pytorch/strategies/strategy.py", line 239, in optimizer_step
    return self.precision_plugin.optimizer_step(optimizer, model=model, closure=closure, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/lightning/pytorch/plugins/precision/amp.py", line 76, in optimizer_step
    return super().optimizer_step(optimizer, model=model, closure=closure, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/lightning/pytorch/plugins/precision/precision.py", line 123, in optimizer_step
    return optimizer.step(closure=closure, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torch/optim/lr_scheduler.py", line 133, in wrapper
    return func.__get__(opt, opt.__class__)(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torch/optim/optimizer.py", line 516, in wrapper
    out = func(*args, **kwargs)
          ^^^^^^^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torch/utils/_contextlib.py", line 120, in decorate_context
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/bitsandbytes/optim/optimizer.py", line 272, in step
    loss = closure()
           ^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/lightning/pytorch/plugins/precision/precision.py", line 109, in _wrap_closure
    closure_result = closure()
                     ^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/lightning/pytorch/loops/optimization/automatic.py", line 146, in __call__
    self._result = self.closure(*args, **kwargs)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torch/utils/_contextlib.py", line 120, in decorate_context
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/lightning/pytorch/loops/optimization/automatic.py", line 140, in closure
    self._backward_fn(step_output.closure_loss)
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/lightning/pytorch/loops/optimization/automatic.py", line 241, in backward_fn
    call._call_strategy_hook(self.trainer, "backward", loss, optimizer)
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/lightning/pytorch/trainer/call.py", line 328, in _call_strategy_hook
    output = fn(*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/lightning/pytorch/strategies/strategy.py", line 213, in backward
    self.precision_plugin.backward(closure_loss, self.lightning_module, optimizer, *args, **kwargs)
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/lightning/pytorch/plugins/precision/precision.py", line 73, in backward
    model.backward(tensor, *args, **kwargs)
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/lightning/pytorch/core/module.py", line 1097, in backward
    loss.backward(*args, **kwargs)
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torch/_tensor.py", line 647, in backward
    torch.autograd.backward(
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torch/autograd/__init__.py", line 354, in backward
    _engine_run_backward(
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torch/autograd/graph.py", line 829, in _engine_run_backward
    return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torch/autograd/function.py", line 311, in apply
    return user_fn(self, *args)
           ^^^^^^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torchao/float8/float8_linear.py", line 197, in backward
    grad_weight = torch.mm(
                  ^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torchao/float8/float8_training_tensor.py", line 374, in __torch_dispatch__
    return FLOAT8_OPS_TABLE[func](func, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torchao/float8/float8_ops.py", line 385, in float8_mm
    tensor_out = addmm_float8_unwrapped(
                 ^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user_1/miniconda3/envs/deneme/lib/python3.11/site-packages/torchao/float8/float8_ops.py", line 72, in addmm_float8_unwrapped
    output = torch._scaled_mm(
             ^^^^^^^^^^^^^^^^^
RuntimeError: Expected trailing dimension of mat1 to be divisible by 16 but got mat1 shape: (6144x1).
Epoch 0:   0%|          | 0/1000 [00:09<?, ?it/s] 

I couldn't turn off checkpointing since my VRAM is not sufficient to train the model without gradient checkpointing.

The error is interesting since I am also applying the filtering function of:

def module_filter_fn(mod: torch.nn.Module, fqn: str):
    # don't convert linear modules with weight dimensions not divisible by 16
    if isinstance(mod, torch.nn.Linear):
        if mod.in_features % 16 != 0 or mod.out_features % 16 != 0:
            return False
    else:
        return False
    
    return True

[vkuzo]

@faruknane , share a printout of your model with float8 applied? It looks like there is a 6144x1 tensor flowing into the backward of one of your linears, which is not supported. If you find out which one, TBH just skip float8 from that linear? I'd suspect last layer of your model without any other info.