Fix compiler toolkit CI by removing duplicated buffer registration

[yiming0416]

self.rope.cache and self.freqs_cis are the same tensor object registered as buffers on two different modules. Tracing would see them as two distinct graph inputs for the same underlying data.

This PR removes the register_buffer from RoPE and just store cache as a plain tensor attribute there, keeping only the Decoder-level register_buffer("freqs_cis", ...).

This fixes the compiler toolkit CI:

NGPU=4 TRAIN_FILE=torchtitan.experiments.compiler_toolkit.train MODULE=compiler_toolkit.llama3 CONFIG=compiler_toolkit_llama3_debugmodel ./run_train.sh --parallelism.data_parallel_shard_degree=2 --parallelism.tensor_parallel_degree=2

[yiming0416]

@tianyu-l This fixes the compiler toolkit CI failure after your config system change.

let me know if you are okay with changing the model code. Otherwise I can fix it from the compiler toolkit experiment.

[tianyu-l]

sounds OK to me, @fegin please take a look

[wconstab]

This PR removes the register_buffer from RoPE and just store cache as a plain tensor attribute there
this doesn't have any impact to checkpointing does it?

[yiming0416]

This PR removes the register_buffer from RoPE and just store cache as a plain tensor attribute there

this doesn't have any impact to checkpointing does it?

@wconstab Originally the cache was registered as a non-persistent buffer, which won't appear in state_dict, so I assume it shouldn't affect checkpointing?

[yiming0416]

@fegin could you take a look? thanks!

[fegin]

LGTM, I don't think checkpointing is a problem.

[wconstab]

one other consideration, iiuc when registered as a non-persisten buffer, 'cache' will at least be 'moved' to device when the module is moved .to(device). Having it as a plain tensor will not do this. Will this break our initialization flow (starting as 'meta' and moving to 'cuda' for example?)