support skip atten in export

Summary:
Support export for llama model variants with attention layer skipping. We only need to specify the attention skip patterns in config.json in layer_type. E.g., 

"layer_types": [
    "full_attention",
    "full_attention",
    "full_attention",
    "skip_attention",
    "skip_attention",
    "skip_attention"
  ]

Differential Revision: D88399533

Author: Hanxian97

examples/models/llama/llama_transformer.py

@@ -132,6 +132,9 @@ def forward(self, x, freqs_cos, freqs_sin, attn_options: ForwardOptions):  # x:
         )
         if not isinstance(self.attention, AttentionSkip):
             h = x + h
+        else:
+            h = x
+            attn_options_update = None
 
         if hasattr(self, "block_sparse_moe"):
             out = h + self.block_sparse_moe(self.ffn_norm(h))
@@ -272,6 +275,10 @@ def construct_transformer(model_args: ModelArgs) -> Transformer:
                     norm_eps=model_args.norm_eps,
                 )
             )
+        elif model_args.layer_types and model_args.layer_types[layer_id] == "skip_attention":
+            attention = AttentionSkip()
+            transformer_block = TransformerBlock(model_args, attention)
+            layers.append(transformer_block)
         else:
             attention = cls(
                 model_args, layer_id, rope, **model_args.attention_kwargs