fix typo in auto wrap policy

Fix a typo. The FSDP wrapper should wrap the `MllamaCrossAttentionDecoderLayer`, which was missing.
meta-llama · Nov 19, 2024 · a62aff3 · a62aff3
1 parent 4250439
commit a62aff3
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/llama_recipes/finetuning.py b/src/llama_recipes/finetuning.py
@@ -237,13 +237,13 @@ def main(**kwargs):
             freeze_transformer_layers(model, train_config.num_freeze_layers)
 
         mixed_precision_policy, wrapping_policy = get_policies(fsdp_config, rank)
-        # Create the FSDP wrapper for MllamaSelfAttentionDecoderLayer,MllamaSelfAttentionDecoderLayer,MllamaVisionEncoderLayer in vision models
+        # Create the FSDP wrapper for MllamaSelfAttentionDecoderLayer,MllamaCrossAttentionDecoderLayer,MllamaVisionEncoderLayer in vision models
         if is_vision:
             my_auto_wrapping_policy = fsdp_auto_wrap_policy(
                 model,
                 [
                     MllamaSelfAttentionDecoderLayer,
-                    MllamaSelfAttentionDecoderLayer,
+                    MllamaCrossAttentionDecoderLayer,
                     MllamaVisionEncoderLayer,
                 ],
             )