Fix rope source transformation error message (#5630)

Summary:
Pull Request resolved: #5630

The error message here accessed attention incorrectly.

Reviewed By: RylanC24, larryliu0820

Differential Revision: D63394466

fbshipit-source-id: 6e3d859ba20f27c26fc44b53630fbcd0eb947561

Author: swolchok

examples/models/llama2/source_transformation/rope.py

@@ -16,11 +16,11 @@ def materialze_broadcast_of_rope_freq_cis(
     assert module.freqs_cos.dim() == 2
     dim0 = module.freqs_cos.size(0)
     dim1 = module.freqs_cos.size(1)
+    module_attention = module.layers[0].attention
     assert (
-        module.layers[0].attention.n_local_kv_heads
-        == module.layers[0].attention.n_local_heads
-    ), f"For rope freqs to be materialzed for broadcast q, k, v num heads must match. For q got {module.attention.n_kv_heads} for k got {module.attention.n_local_heads} and v got {module.attention.n_local_kv_heads}"
-    num_heads = module.layers[0].attention.n_local_heads
+        module_attention.n_local_kv_heads == module_attention.n_local_heads
+    ), f"For rope freqs to be materialized for broadcast, q, k, v num heads must match. For q got {module_attention.n_kv_heads} for k got {module_attention.n_local_heads} and v got {module_attention.n_local_kv_heads}"
+    num_heads = module_attention.n_local_heads
     module.freqs_cos = module.freqs_cos.view(dim0, 1, dim1)
     module.freqs_cos = module.freqs_cos.expand(dim0, num_heads, dim1).contiguous()
     assert module.freqs_sin.dim() == 2