Checkpointing is not compatible with .grad(), please use .backward() if possible

[chengmengli06]

when training deepseek v3 model using cuda_graph, it reports the following error:
Checkpointing is not compatible with .grad(), please use .backward() if possible

setting the following args:
--recompute-granularity selective
--recompute-modules mla_up_proj mlp

--cuda-graph-impl transformer_engine
--cuda-graph-scope full

[CodersAcademy006]

To fix this, replace the functional torch.autograd.grad(loss, inputs, ...) call with the standard loss.backward(create_graph=True, retain_graph=True). This triggers the full autograd engine, which correctly manages the re-entrant checkpointing hooks. You can then manually collect the gradients from the .grad attributes of your inputs (e.g., grads = [p.grad for p in inputs]). If you strictly require .grad(), you must temporarily disable activation recomputation (model.config.recompute_granularity = None) for that specific forward pass, provided you have the GPU memory to spare.

Hope this helps!!