[Usage]: torch.func.jvp不支持，有什么替代方案？-pytorch-AtomGit

已关闭

[Usage]: torch.func.jvp不支持，有什么替代方案？ #1498

hw_yhzhen创建于 1月8日关闭于 1月20日

1月8日创建

在提交新问题之前，请确保您已经在社区中搜索过相关问题，并使用了社区中提供的资源/工具后，仍未找到满意的解决方式。

环境信息

使用torch.func.jvp api不支持，有什么替代方案？
报错如下:
XXXXXXXXXXXXXX.py", line 626, in meanflow_forward
[rank1]: (u, (prefix_out, pad_masks)), dudt = jvp(
[rank1]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/torch/_functorch/eager_transforms.py", line 1042, in jvp
[rank1]: return _jvp_with_argnums(
[rank1]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/torch/_functorch/eager_transforms.py", line 1094, in _jvp_with_argnums
[rank1]: flat_duals = tuple(
[rank1]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/torch/_functorch/eager_transforms.py", line 1095, in
[rank1]: fwAD.make_dual(p, t) for p, t in zip(flat_primals, flat_tangents)
[rank1]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/torch/autograd/forward_ad.py", line 110, in make_dual
[rank1]: _maybe_load_decompositions()
[rank1]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/torch/autograd/forward_ad.py", line 67, in _maybe_load_decompositions
[rank1]: from torch._decomp import decompositions_for_jvp # noqa: F401
[rank1]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/torch/_decomp/decompositions_for_jvp.py", line 326, in
[rank1]: _register_jit_decomposition_for_jvp(torch.ops.aten.nll_loss_backward.default)
[rank1]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/torch/_decomp/decompositions_for_jvp.py", line 93, in _register_jit_decomposition_for_jvp
[rank1]: graph = torch.jit.script(decomp_fn).graph
[rank1]: AttributeError: 'function' object has no attribute 'graph'
[ERROR] 2026-01-08-15:14:16 (PID:3806, Device:1, RankID:-1) ERR99999 UNKNOWN applicaiton exception
/usr/local/python3.10.18/lib/python3.10/site-packages/torch_npu/contrib/transfer_to_npu.py:161: UserWarning: Cannot create tensor with interal format while allow_internel_format=False, tensor will be created with base format. (Triggered internally at build/CMakeFiles/torch_npu.dir/compiler_depend.ts:335.)
return fn(*args, **kwargs)
[rank0]: Traceback (most recent call last):
[rank0]: File "/data/private/XX_add/./training/train_XX.py", line 131, in
[rank0]: main(all_config=config)
[rank0]: File "/data/private/XX_add/./training/train_XX.py", line 105, in main
[rank0]: trainer.train(resume_from_checkpoint=all_config['training_args'].resume_from_checkpoint)
[rank0]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/transformers/trainer.py", line 2171, in train
[rank0]: return inner_training_loop(
[rank0]: File "XXXXXXXXX_trainer.py", line 652, in _inner_training_loop
[rank0]: tr_loss_step, all_loss = self.training_step(model, inputs) # modified,return all_loss
[rank0]: File "XXXXXXXXX_trainer.py", line 294, in training_step
[rank0]: loss = self.compute_loss(model, inputs, return_outputs=False) # change return_outputs to True
[rank0]: File "XXXXXXXXX_trainer.py", line 1082, in compute_loss
[rank0]: outputs = model(inputs)
[rank0]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
[rank0]: return self._call_impl(*args, **kwargs)
[rank0]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl
[rank0]: return forward_call(*args, **kwargs)
[rank0]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/deepspeed/utils/nvtx.py", line 20, in wrapped_fn
[rank0]: ret_val = func(*args, **kwargs)
[rank0]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/deepspeed/runtime/engine.py", line 2105, in forward
[rank0]: loss = self.module(*inputs, **kwargs)
[rank0]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
[rank0]: return self._call_impl(*args, **kwargs)
[rank0]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1845, in _call_impl
[rank0]: return inner()
[rank0]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1793, in inner
[rank0]: result = forward_call(*args, **kwargs)
[rank0]: File "XXXXXXXXXXXXXX.py", line 223, in forward
[rank0]: losses, loss_lm = self.model.forward(images, img_masks, lang_tokens, lang_masks, state, actions, noise, time, labels)
[rank0]: File "XXXXXXXXXXXXXX.py", line 496, in forward
[rank0]: loss_action, loss_lm = self.meanflow_forward(
[rank0]: File "XXXXXXXXXXXXXX.py", line 626, in meanflow_forward
[rank0]: (u, (prefix_out, pad_masks)), dudt = jvp(
[rank0]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/torch/_functorch/eager_transforms.py", line 1042, in jvp
[rank0]: return _jvp_with_argnums(
[rank0]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/torch/_functorch/eager_transforms.py", line 1094, in _jvp_with_argnums
[rank0]: flat_duals = tuple(
[rank0]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/torch/_functorch/eager_transforms.py", line 1095, in
[rank0]: fwAD.make_dual(p, t) for p, t in zip(flat_primals, flat_tangents)
[rank0]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/torch/autograd/forward_ad.py", line 110, in make_dual
[rank0]: _maybe_load_decompositions()
[rank0]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/torch/autograd/forward_ad.py", line 67, in _maybe_load_decompositions
[rank0]: from torch._decomp import decompositions_for_jvp # noqa: F401
[rank0]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/torch/_decomp/decompositions_for_jvp.py", line 326, in
[rank0]: _register_jit_decomposition_for_jvp(torch.ops.aten.nll_loss_backward.default)
[rank0]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/torch/_decomp/decompositions_for_jvp.py", line 93, in _register_jit_decomposition_for_jvp
[rank0]: graph = torch.jit.script(decomp_fn).graph
[rank0]: AttributeError: 'function' object has no attribute 'graph'
[ERROR] 2026-01-08-15:14:19 (PID:3805, Device:0, RankID:-1) ERR99999 UNKNOWN applicaiton exception

使用场景及问题

欢迎加入社区，感谢您对社区的贡献 🎉!

kkjocker成员

1月8日关联了pull request：fix transfer_to_npu close torch.jit.script in every situation

haiyan8

1月8日评论：

在使用过程中是否用了rom torch_npu.contrib import transfer_to_npu 。如果用了，先去掉这个，再测下。

kkjocker成员

1月8日关联了pull request：fix transfer_to_npu close torch.jit.script in every situation

hw_yhzhen

1月9日评论：

[rank1]: Traceback (most recent call last):
[rank1]: File "XXXXX", line 131, in
[rank1]: main(all_config=config)
[rank1]: File "XXXXX", line 105, in main
[rank1]: trainer.train(resume_from_checkpoint=all_config['training_args'].resume_from_checkpoint)
[rank1]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/transformers/trainer.py", line 2171, in train
[rank1]: return inner_training_loop(
[rank1]: File "XXXXX", line 652, in _inner_training_loop
[rank1]: tr_loss_step, all_loss = self.training_step(model, inputs) # modified,return all_loss
[rank1]: File "XXXXX", line 294, in training_step
[rank1]: loss = self.compute_loss(model, inputs, return_outputs=False) # change return_outputs to True
[rank1]: File "XXXXX", line 1082, in compute_loss
[rank1]: outputs = model(inputs)
[rank1]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
[rank1]: return self._call_impl(*args, **kwargs)
[rank1]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl
[rank1]: return forward_call(*args, **kwargs)
[rank1]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/deepspeed/utils/nvtx.py", line 20, in wrapped_fn
[rank1]: ret_val = func(*args, **kwargs)
[rank1]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/deepspeed/runtime/engine.py", line 2105, in forward
[rank1]: loss = self.module(*inputs, **kwargs)
[rank1]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
[rank1]: return self._call_impl(*args, **kwargs)
[rank1]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1845, in _call_impl
[rank1]: return inner()
[rank1]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1793, in inner
[rank1]: result = forward_call(*args, **kwargs)
[rank1]: File "XXXXX.py", line 223, in forward
[rank1]: losses, loss_lm = self.model.forward(images, img_masks, lang_tokens, lang_masks, state, actions, noise, time, labels)
[rank1]: File "XXXXX.py", line 496, in forward
[rank1]: loss_action, loss_lm = self.meanflow_forward(
[rank1]: File "XXXXX.py", line 626, in meanflow_forward
[rank1]: (u, (prefix_out, pad_masks)), dudt = jvp(
[rank1]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/torch/_functorch/eager_transforms.py", line 1042, in jvp
[rank1]: return _jvp_with_argnums(
[rank1]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/torch/_functorch/eager_transforms.py", line 1101, in _jvp_with_argnums
[rank1]: result_duals = func(*duals)
[rank1]: File "XXXXX.py", line 614, in fn_for_jvp
[rank1]: u, prefix_out, pad_masks = self._compute_u_for_jvp(
[rank1]: File "XXXXX.py", line 532, in _compute_u_for_jvp
[rank1]: suffix_embs, suffix_pad_masks, suffix_att_masks = self.embed_suffix(state, z, t)
[rank1]: File "XXXXX.py", line 466, in embed_suffix
[rank1]: action_emb = self.action_in_proj(noisy_actions)
[rank1]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
[rank1]: return self._call_impl(*args, **kwargs)
[rank1]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl
[rank1]: return forward_call(*args, **kwargs)
[rank1]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/torch/nn/modules/linear.py", line 125, in forward
[rank1]: return F.linear(input, self.weight, self.bias)
[rank1]: RuntimeError: wrapper == nullptr INTERNAL ASSERT FAILED at "/pytorch/aten/src/ATen/functorch/Interpreter.cpp":97, please report a bug to PyTorch.
/usr/local/python3.10.18/lib/python3.10/site-packages/torch_npu/contrib/transfer_to_npu.py:161: UserWarning: Cannot create tensor with interal format while allow_internel_format=False, tensor will be created with base format. (Triggered internally at build/CMakeFiles/torch_npu.dir/compiler_depend.ts:335.)
return fn(*args, **kwargs)
/usr/local/python3.10.18/lib/python3.10/site-packages/torch_npu/contrib/transfer_to_npu.py:252: RuntimeWarning: torch.jit.script failed, failing back to eager function
warnings.warn(
[rank0]:[W109 14:47:54.144176216 compiler_depend.ts:250] Warning: CAUTION: The operator 'aten::_efficientzerotensor' is not currently supported on the NPU backend and will fall back to run on the CPU. This may have performance implications. (function npu_cpu_fallback)
[ERROR] 2026-01-09-14:47:56 (PID:20477, Device:1, RankID:-1) ERR99999 UNKNOWN applicaiton exception

hw_yhzhen

1月9日评论：

用你们提供的快速规避方案报错如上

hw_yhzhen

1月9日评论：

35.)
tangent_t = torch.ones_like(t) # shape: (bsize,)
[rank1]: Traceback (most recent call last):
[rank1]: File "XXXXX", line 132, in
[rank1]: main(all_config=config)
[rank1]: File "XXXXX", line 106, in main
[rank1]: trainer.train(resume_from_checkpoint=all_config['training_args'].resume_from_checkpoint)
[rank1]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/transformers/trainer.py", line 2171, in train
[rank1]: return inner_training_loop(
[rank1]: File "XXXXX", line 653, in _inner_training_loop
[rank1]: tr_loss_step, all_loss = self.training_step(model, inputs) # modified,return all_loss
[rank1]: File "XXXXX", line 295, in training_step
[rank1]: loss = self.compute_loss(model, inputs, return_outputs=False) # change return_outputs to True
[rank1]: File "XXXXX", line 1083, in compute_loss
[rank1]: outputs = model(inputs)
[rank1]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
[rank1]: return self._call_impl(*args, **kwargs)
[rank1]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl
[rank1]: return forward_call(*args, **kwargs)
[rank1]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/deepspeed/utils/nvtx.py", line 20, in wrapped_fn
[rank1]: ret_val = func(*args, **kwargs)
[rank1]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/deepspeed/runtime/engine.py", line 2105, in forward
[rank1]: loss = self.module(*inputs, **kwargs)
[rank1]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
[rank1]: return self._call_impl(*args, **kwargs)
[rank1]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1845, in _call_impl
[rank1]: return inner()
[rank1]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1793, in inner
[rank1]: result = forward_call(*args, **kwargs)
[rank1]: File "XXXXX", line 224, in forward
[rank1]: losses, loss_lm = self.model.forward(images, img_masks, lang_tokens, lang_masks, state, actions, noise, time, labels)
[rank1]: File "XXXXX", line 497, in forward
[rank1]: loss_action, loss_lm = self.meanflow_forward(
[rank1]: File "XXXXX", line 627, in meanflow_forward
[rank1]: (u, (prefix_out, pad_masks)), dudt = jvp(
[rank1]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/torch/_functorch/eager_transforms.py", line 1042, in jvp
[rank1]: return _jvp_with_argnums(
[rank1]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/torch/_functorch/eager_transforms.py", line 1101, in _jvp_with_argnums
[rank1]: result_duals = func(*duals)
[rank1]: File "XXXXX", line 615, in fn_for_jvp
[rank1]: u, prefix_out, pad_masks = self._compute_u_for_jvp(
[rank1]: File "XXXXX", line 533, in _compute_u_for_jvp
[rank1]: suffix_embs, suffix_pad_masks, suffix_att_masks = self.embed_suffix(state, z, t)
[rank1]: File "XXXXX", line 467, in embed_suffix
[rank1]: action_emb = self.action_in_proj(noisy_actions)
[rank1]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
[rank1]: return self._call_impl(*args, **kwargs)
[rank1]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl
[rank1]: return forward_call(*args, **kwargs)
[rank1]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/torch/nn/modules/linear.py", line 125, in forward
[rank1]: return F.linear(input, self.weight, self.bias)
[rank1]: RuntimeError: wrapper == nullptr INTERNAL ASSERT FAILED at "/pytorch/aten/src/ATen/functorch/Interpreter.cpp":97, please report a bug to PyTorch.
XXXXX:623: UserWarning: Cannot create tensor with interal format while allow_internel_format=False, tensor will be created with base format. (Triggered internally at build/CMakeFiles/torch_npu.dir/compiler_depend.ts:335.)
tangent_t = torch.ones_like(t) # shape: (bsize,)
[rank0]:[W109 15:22:53.614385807 compiler_depend.ts:250] Warning: CAUTION: The operator 'aten::_efficientzerotensor' is not currently supported on the NPU backend and will fall back to run on the CPU. This may have performance implications. (function npu_cpu_fallback)
[rank0]: Traceback (most recent call last):
[rank0]: File "XXXXX", line 132, in
[rank0]: main(all_config=config)
[rank0]: File "XXXXX", line 106, in main
[rank0]: trainer.train(resume_from_checkpoint=all_config['training_args'].resume_from_checkpoint)
[rank0]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/transformers/trainer.py", line 2171, in train
[rank0]: return inner_training_loop(
[rank0]: File "XXXXX", line 653, in _inner_training_loop
[rank0]: tr_loss_step, all_loss = self.training_step(model, inputs) # modified,return all_loss
[rank0]: File "XXXXX", line 295, in training_step
[rank0]: loss = self.compute_loss(model, inputs, return_outputs=False) # change return_outputs to True
[rank0]: File "XXXXX", line 1083, in compute_loss
[rank0]: outputs = model(inputs)
[rank0]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
[rank0]: return self._call_impl(*args, **kwargs)
[rank0]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl
[rank0]: return forward_call(*args, **kwargs)
[rank0]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/deepspeed/utils/nvtx.py", line 20, in wrapped_fn
[rank0]: ret_val = func(*args, **kwargs)
[rank0]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/deepspeed/runtime/engine.py", line 2105, in forward
[rank0]: loss = self.module(*inputs, **kwargs)
[rank0]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
[rank0]: return self._call_impl(*args, **kwargs)
[rank0]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1845, in _call_impl
[rank0]: return inner()
[rank0]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1793, in inner
[rank0]: result = forward_call(*args, **kwargs)
[rank0]: File "XXXXX", line 224, in forward
[rank0]: losses, loss_lm = self.model.forward(images, img_masks, lang_tokens, lang_masks, state, actions, noise, time, labels)
[rank0]: File "XXXXX", line 497, in forward
[rank0]: loss_action, loss_lm = self.meanflow_forward(
[rank0]: File "XXXXX", line 627, in meanflow_forward
[rank0]: (u, (prefix_out, pad_masks)), dudt = jvp(
[rank0]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/torch/_functorch/eager_transforms.py", line 1042, in jvp
[rank0]: return _jvp_with_argnums(
[rank0]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/torch/_functorch/eager_transforms.py", line 1101, in _jvp_with_argnums
[rank0]: result_duals = func(*duals)
[rank0]: File "XXXXX", line 615, in fn_for_jvp
[rank0]: u, prefix_out, pad_masks = self._compute_u_for_jvp(
[rank0]: File "XXXXX", line 533, in _compute_u_for_jvp
[rank0]: suffix_embs, suffix_pad_masks, suffix_att_masks = self.embed_suffix(state, z, t)
[rank0]: File "XXXXX", line 467, in embed_suffix
[rank0]: action_emb = self.action_in_proj(noisy_actions)
[rank0]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
[rank0]: return self._call_impl(*args, **kwargs)
[rank0]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl
[rank0]: return forward_call(*args, **kwargs)
[rank0]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/torch/nn/modules/linear.py", line 125, in forward
[rank0]: return F.linear(input, self.weight, self.bias)
[rank0]: RuntimeError: wrapper == nullptr INTERNAL ASSERT FAILED at "/pytorch/aten/src/ATen/functorch/Interpreter.cpp":97, please report a bug to PyTorch.
[ERROR] 2026-01-09-15:22:55 (PID:25544, Device:1, RankID:-1) ERR99999 UNKNOWN applicaiton exception
[ERROR] 2026-01-09-15:22:55 (PID:25543, Device:0, RankID:-1) ERR99999 UNKNOWN applicaiton exception

hw_yhzhen

1月9日评论：

去除transfer_to_npu，使用工具迁移，报错如上

haiyan8

1月12日评论：

你的torch是什么版本的。代码能在gpu上跑通吗

hw_yhzhen

1月13日评论：

action_in_proj = nn.Linear(32, 1024)

def _compute_u_for_jvp(z_input, t_input):
# embs, pad_masks, att_masks = embed_suffix(state, z, t)
nosiy_actions = z_input.to(dtype=torch.float32)
action_in_proj.to(dtype=torch.float32)
action_emb = action_in_proj(nosiy_actions)

def fn_for_jvp(z_input, t_input):
u, out, masks = _compute_u_for_jvp(z_input, t_input)

z = torch.randn(6, 50, 32).npu()
t = torch.randn(6).npu()
tangent_z = torch.randn(6,50,32).npu()
tangent_t = torch.ones(6).npu()

(u, (out, mask)), dudt = func.jvp(fn_for_jvp, (z,t), (tangent_z, tangent_t), has_aux=True)

hw_yhzhen

1月13日评论：

import torch
import torch_npu
import torch.func as func # jvp 核心库
import torch.nn as nn # nn.Linear 核心库

action_in_proj = nn.Linear(32, 1024).npu()

@torch.func.functionalize
def _compute_u_for_jvp(z_input, t_input):
nosiy_actions = z_input.to(dtype=torch.float32)
action_in_proj.to(dtype=torch.float32)
action_emb = action_in_proj(nosiy_actions)

u = action_emb
out = torch.mean(action_emb, dim=-1)
masks = torch.ones_like(out)

# 返回3个值，满足fn_for_jvp的解包要求
return u, out, masks

def fn_for_jvp(z_input, t_input):
u, out, masks = _compute_u_for_jvp(z_input, t_input)

return u, (out, masks)

z = torch.randn(6, 50, 32).npu()
t = torch.randn(6).npu()
tangent_z = torch.randn(6,50,32).npu()
tangent_t = torch.ones(6).npu()

output = func.jvp(fn_for_jvp, (z,t), (tangent_z, tangent_t), has_aux=True)
print(output)

hw_yhzhen

1月14日评论：

[rank1]: Traceback (most recent call last):
[rank1]: File "/data/private/mirro_add_msft_bak/./training/train_mirro.py", line 132, in
[rank1]: main(all_config=config)
[rank1]: File "/data/private/mirro_add_msft_bak/./training/train_mirro.py", line 106, in main
[rank1]: trainer.train(resume_from_checkpoint=all_config['training_args'].resume_from_checkpoint)
[rank1]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/transformers/trainer.py", line 2171, in train
[rank1]: return inner_training_loop(
[rank1]: File "/data/private/mirro_add_msft_bak/models/mirro_trainer.py", line 653, in _inner_training_loop
[rank1]: tr_loss_step, all_loss = self.training_step(model, inputs) # modified,return all_loss
[rank1]: File "/data/private/mirro_add_msft_bak/models/mirro_trainer.py", line 295, in training_step
[rank1]: loss = self.compute_loss(model, inputs, return_outputs=False) # change return_outputs to True
[rank1]: File "/data/private/mirro_add_msft_bak/models/mirro_trainer.py", line 1083, in compute_loss
[rank1]: outputs = model(inputs)
[rank1]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
[rank1]: return self._call_impl(*args, **kwargs)
[rank1]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl
[rank1]: return forward_call(*args, **kwargs)
[rank1]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/deepspeed/utils/nvtx.py", line 20, in wrapped_fn
[rank1]: ret_val = func(*args, **kwargs)
[rank1]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/deepspeed/runtime/engine.py", line 2105, in forward
[rank1]: loss = self.module(*inputs, **kwargs)
[rank1]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
[rank1]: return self._call_impl(*args, **kwargs)
[rank1]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1845, in _call_impl
[rank1]: return inner()
[rank1]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1793, in inner
[rank1]: result = forward_call(*args, **kwargs)
[rank1]: File "/data/private/mirro_add_msft_bak/models/mirro/modeling_mirro_meanflow.py", line 224, in forward
[rank1]: losses, loss_lm = self.model.forward(images, img_masks, lang_tokens, lang_masks, state, actions, noise, time, labels)
[rank1]: File "/data/private/mirro_add_msft_bak/models/mirro/modeling_mirro_meanflow.py", line 497, in forward
[rank1]: loss_action, loss_lm = self.meanflow_forward(
[rank1]: File "/data/private/mirro_add_msft_bak/models/mirro/modeling_mirro_meanflow.py", line 627, in meanflow_forward
[rank1]: (u, (prefix_out, pad_masks)), dudt = jvp(
[rank1]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/torch/_functorch/eager_transforms.py", line 1042, in jvp
[rank1]: return _jvp_with_argnums(
[rank1]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/torch/_functorch/eager_transforms.py", line 1101, in _jvp_with_argnums
[rank1]: result_duals = func(*duals)
[rank1]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/torch/_functorch/eager_transforms.py", line 1645, in wrapped
[rank1]: func_outputs = func(*func_args, **func_kwargs)
[rank1]: File "/data/private/mirro_add_msft_bak/models/mirro/modeling_mirro_meanflow.py", line 615, in fn_for_jvp
[rank1]: u, prefix_out, pad_masks = self._compute_u_for_jvp(
[rank1]: File "/data/private/mirro_add_msft_bak/models/mirro/modeling_mirro_meanflow.py", line 541, in _compute_u_for_jvp
[rank1]: (prefix_out, suffix_out), _ = self.paligemma_with_expert.forward(
[rank1]: File "/data/private/mirro_add_msft_bak/models/mirro/paligemma_with_expert.py", line 283, in forward
[rank1]: query_states = apply_rope(query_states, position_ids)
[rank1]: File "/data/private/mirro_add_msft_bak/models/mirro/paligemma_with_expert.py", line 43, in apply_rope
[rank1]: res[..., :d_half] = x1 * cos - x2 * sin
[rank1]: NotImplementedError: Trying to use forward AD with aten::copy that does not support it.
/data/private/mirro_add_msft_bak/models/mirro/modeling_mirro_meanflow.py:623: UserWarning: Cannot create tensor with interal format while allow_internel_format=False, tensor will be created with base format. (Triggered internally at build/CMakeFiles/torch_npu.dir/compiler_depend.ts:335.)
tangent_t = torch.ones_like(t) # shape: (bsize,)
[rank0]:[W114 18:13:19.853574886 compiler_depend.ts:250] Warning: CAUTION: The operator 'aten::_efficientzerotensor' is not currently supported on the NPU backend and will fall back to run on the CPU. This may have performance implications. (function npu_cpu_fallback)
[rank0]: Traceback (most recent call last):
[rank0]: File "/data/private/mirro_add_msft_bak/./training/train_mirro.py", line 132, in
[rank0]: main(all_config=config)
[rank0]: File "/data/private/mirro_add_msft_bak/./training/train_mirro.py", line 106, in main
[rank0]: trainer.train(resume_from_checkpoint=all_config['training_args'].resume_from_checkpoint)
[rank0]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/transformers/trainer.py", line 2171, in train
[rank0]: return inner_training_loop(
[rank0]: File "/data/private/mirro_add_msft_bak/models/mirro_trainer.py", line 653, in _inner_training_loop
[rank0]: tr_loss_step, all_loss = self.training_step(model, inputs) # modified,return all_loss
[rank0]: File "/data/private/mirro_add_msft_bak/models/mirro_trainer.py", line 295, in training_step
[rank0]: loss = self.compute_loss(model, inputs, return_outputs=False) # change return_outputs to True
[rank0]: File "/data/private/mirro_add_msft_bak/models/mirro_trainer.py", line 1083, in compute_loss
[rank0]: outputs = model(inputs)
[rank0]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
[rank0]: return self._call_impl(*args, **kwargs)
[rank0]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl
[rank0]: return forward_call(*args, **kwargs)
[rank0]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/deepspeed/utils/nvtx.py", line 20, in wrapped_fn
[rank0]: ret_val = func(*args, **kwargs)
[rank0]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/deepspeed/runtime/engine.py", line 2105, in forward
[rank0]: loss = self.module(*inputs, **kwargs)
[rank0]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
[rank0]: return self._call_impl(*args, **kwargs)
[rank0]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1845, in _call_impl
[rank0]: return inner()
[rank0]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1793, in inner
[rank0]: result = forward_call(*args, **kwargs)
[rank0]: File "/data/private/mirro_add_msft_bak/models/mirro/modeling_mirro_meanflow.py", line 224, in forward
[rank0]: losses, loss_lm = self.model.forward(images, img_masks, lang_tokens, lang_masks, state, actions, noise, time, labels)
[rank0]: File "/data/private/mirro_add_msft_bak/models/mirro/modeling_mirro_meanflow.py", line 497, in forward
[rank0]: loss_action, loss_lm = self.meanflow_forward(
[rank0]: File "/data/private/mirro_add_msft_bak/models/mirro/modeling_mirro_meanflow.py", line 627, in meanflow_forward
[rank0]: (u, (prefix_out, pad_masks)), dudt = jvp(
[rank0]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/torch/_functorch/eager_transforms.py", line 1042, in jvp
[rank0]: return _jvp_with_argnums(
[rank0]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/torch/_functorch/eager_transforms.py", line 1101, in _jvp_with_argnums
[rank0]: result_duals = func(*duals)
[rank0]: File "/usr/local/python3.10.18/lib/python3.10/site-packages/torch/_functorch/eager_transforms.py", line 1645, in wrapped
[rank0]: func_outputs = func(*func_args, **func_kwargs)
[rank0]: File "/data/private/mirro_add_msft_bak/models/mirro/modeling_mirro_meanflow.py", line 615, in fn_for_jvp
[rank0]: u, prefix_out, pad_masks = self._compute_u_for_jvp(
[rank0]: File "/data/private/mirro_add_msft_bak/models/mirro/modeling_mirro_meanflow.py", line 541, in _compute_u_for_jvp
[rank0]: (prefix_out, suffix_out), _ = self.paligemma_with_expert.forward(
[rank0]: File "/data/private/mirro_add_msft_bak/models/mirro/paligemma_with_expert.py", line 283, in forward
[rank0]: query_states = apply_rope(query_states, position_ids)
[rank0]: File "/data/private/mirro_add_msft_bak/models/mirro/paligemma_with_expert.py", line 43, in apply_rope
[rank0]: res[..., :d_half] = x1 * cos - x2 * sin
[rank0]: NotImplementedError: Trying to use forward AD with aten::copy that does not support it.
[ERROR] 2026-01-14-18:13:20 (PID:233245, Device:1, RankID:-1) ERR99999 UNKNOWN applicaiton exception
[ERROR] 2026-01-14-18:13:21 (PID:233244, Device:0, RankID:-1) ERR99999 UNKNOWN applicaiton exception

ascend-robot成员

1月20日关闭了 issue

H hw_yhzhen

1月20日 issue状态由 TODO 改变为 DONE

ascend-robot成员

1月22日添加了label：resolved