NotImplementedError happend, when I trained sd2.1 model with fp8 quantization.

Describe the bug

I train sd2.1 model with fp8 quantization, UserWarning and Errors as folllows:

(1) modelopt/torch/opt/dynamic.py:895: UserWarning: Distributed training is initialized but no parallel_state is set for <class 'modelopt.torch.opt.dynamic.QuantConv2d'>. Using default parallel_state which has data_parallel_group set to the default process group and tensor_parallel_group is unspecified. If you are using tensor parallelism for this module, you should set the parallel_state in its _setup method.
parallelism for this module, you should set the parallel_state in its _setup method.
warnings.warn(
(2) modelopt/torch/quantization/nn/modules/tensor_quantizer.py:1127: UserWarning: Failed to synchronize amax: No backend type associated with
device type cpu, probably because the tensor is on a device which is notsupported by the current distributed backend. This warning can be ignoredif happening during modelopt restore.
warnings.warn(
(3) modelopt/torch/quantization/tensor_quant.py:425: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.detac
h().clone() or sourceTensor.detach().clone().requires_grad_(True), rather than torch.tensor(sourceTensor).
amax = torch.tensor(
(4)
[rank0]: Traceback (most recent call last):
[rank0]: File "/data/train_fp8.py", line 1044, in
[rank0]: main(args)
[rank0]: File "/data/train_fp8.py", line 990, in main
[rank0]: accelerator.backward(loss)
[rank0]: File "/root/miniforge3/envs/torch270/lib/python3.12/site-packages/accelerate/accelerator.py", line 2473, in backward
[rank0]: loss.backward(**kwargs)
[rank0]: File "/root/miniforge3/envs/torch270/lib/python3.12/site-packages/torch/_tensor.py", line 648, in backward
[rank0]: torch.autograd.backward(
[rank0]: File "/root/miniforge3/envs/torch270/lib/python3.12/site-packages/torch/autograd/init.py", line 353, in backward
[rank0]: _engine_run_backward(
[rank0]: File "/root/miniforge3/envs/torch270/lib/python3.12/site-packages/torch/autograd/graph.py", line 824, in _engine_run_backward
[rank0]: return Variable._execution_engine.run_backward( # Calls into the C++ engine to run the backward pass
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/root/miniforge3/envs/torch270/lib/python3.12/site-packages/torch/autograd/function.py", line 307, in apply
[rank0]: return user_fn(self, *args)
[rank0]: ^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/root/miniforge3/envs/torch270/lib/python3.12/site-packages/torch/autograd/function.py", line 431, in backward
[rank0]: raise NotImplementedError(
[rank0]: NotImplementedError: You must implement either the backward or vjp method for your custom autograd.Function to use it with backward mode AD.

Steps/Code to reproduce bug

quant_config = mtq.FP8_DEFAULT_CFG
unet = mtq.quantize(unet, quant_config, forward_loop=None)
unet.train()
unet = accelerator.prepare(unet)

for step, batch in enumerate(train_dataloader):
with accelerator.accumulate([unet, ]):
................
pred = unet(x).sample
loss = （...................）
accelerator.backward(loss)

Expected behavior

System information

• Container used (if applicable):
• OS: Ubuntu 22.04,
• CPU architecture: x86_64
• GPU name : H800
• Number of GPUs: 4
• Library versions (if applicable):
  • Python: 3.12
  • ModelOpt version: 0.31.0
  • CUDA: 12.6
  • PyTorch: 2.7.0
  • Transformers: 4.51.3
  • ONNXRuntime: 1.20.0
  • TensorRT: 10.9.0