KV Cache Initialization Issue During Training

[XiaokunFeng]

Thanks for your great work!

I noticed that during training, the initialization of kv_cache and crossattn_cache is commented out. With the current code, training raises an error because self.crossattn_cache is never defined, while self.kv_cache1 is only initialized to None in _init_.

From my understanding, both self.crossattn_cache and self.kv_cache1 should be explicitly initialized before being used. Could you please clarify the intended behavior here, or suggest how these caches should be properly initialized for training?

Thanks in advance for your help.

[yesiltepe-hidir]

Thanks for letting us know this!

During my experiments, I commented out that part to experiment with with a different cache structures. (not related to infinity-rope, since it is a training-free methodology.)

I updated the code.

I will also double check the rest of the training code.

[XiaokunFeng]

Thanks for your timely reply.

I retrained the model using your updated code, which initializes the KV cache before the training iterations. However, I encountered the following runtime error during training (see details below).

In addition, I noticed that during inference the KV cache supports a maximum of 6 frames, whereas the KV cache is initialized with 21 frames. Should this discrepancy be adjusted or aligned?

I look forward to your response.

[rank0]: Traceback (most recent call last):
[rank0]: File "./infinity-rope/train.py", line 47, in
[rank0]: main()
[rank0]: File "./infinity-rope/train.py", line 41, in main
[rank0]: trainer.train()
[rank0]: File "./infinity-rope/trainer/distillation.py", line 324, in train
[rank0]: extra = self.fwdbwd_one_step(batch, True)
[rank0]: File "./infinity-rope/trainer/distillation.py", line 254, in fwdbwd_one_step
[rank0]: generator_loss.backward()
[rank0]: File "./envs/wan_adjust2/lib/python3.10/site-packages/torch/_tensor.py", line 625, in backward
[rank0]: torch.autograd.backward(
[rank0]: File "./envs/wan_adjust2/lib/python3.10/site-packages/torch/autograd/init.py", line 354, in backward
[rank0]: _engine_run_backward(
[rank0]: File "./envs/wan_adjust2/lib/python3.10/site-packages/torch/autograd/graph.py", line 841, in _engine_run_backward
[rank0]: return Variable._execution_engine.run_backward( # Calls into the C++ engine to run the backward pass
[rank0]: File "./envs/wan_adjust2/lib/python3.10/site-packages/torch/utils/checkpoint.py", line 1158, in unpack_hook
[rank0]: frame.check_recomputed_tensors_match(gid)
[rank0]: File "./envs/wan_adjust2/lib/python3.10/site-packages/torch/utils/checkpoint.py", line 911, in check_recomputed_tensors_match
[rank0]: raise CheckpointError(
[rank0]: torch.utils.checkpoint.CheckpointError: torch.utils.checkpoint: Recomputed values for the following tensors have different metadata than during the forward pass.
[rank0]: tensor at position 24:
[rank0]: saved metadata: {'shape': torch.Size([9360, 1, 64]), 'dtype': torch.complex128, 'device': device(type='cuda', index=0)}
[rank0]: recomputed metadata: {'shape': torch.Size([4680, 1, 64]), 'dtype': torch.complex128, 'device': device(type='cuda', index=0)}
[rank0]: tensor at position 26:
[rank0]: saved metadata: {'shape': torch.Size([9360, 12, 128]), 'dtype': torch.bfloat16, 'device': device(type='cuda', index=0)}
[rank0]: recomputed metadata: {'shape': torch.Size([4680, 12, 128]), 'dtype': torch.bfloat16, 'device': device(type='cuda', index=0)}
[rank0]: tensor at position 27:
[rank0]: saved metadata: {'shape': torch.Size([9360, 12, 128]), 'dtype': torch.bfloat16, 'device': device(type='cuda', index=0)}
[rank0]: recomputed metadata: {'shape': torch.Size([4680, 12, 128]), 'dtype': torch.bfloat16, 'device': device(type='cuda', index=0)}
[rank0]: .

[rank0]: Tip: To see a more detailed error message, either pass debug=True to
[rank0]: torch.utils.checkpoint.checkpoint(...) or wrap the code block
[rank0]: with with torch.utils.checkpoint.set_checkpoint_debug_enabled(True): to
[rank0]: enable checkpoint‑debug mode globally.

[yesiltepe-hidir]

Sure, I will check it out today and make the training code working with Infinity-RoPE