support skip atten in export

jackzhxng · web-flow · commit d859b0119949 · 2025-12-08T23:07:10.000Z
Differential Revision: D88399533 Pull Request resolved: pytorch#16104
diff --git a/examples/models/llama/llama_transformer.py b/examples/models/llama/llama_transformer.py
@@ -272,6 +272,13 @@ 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

Original file line number	Diff line number	Diff line change
`@@ -272,6 +272,13 @@ def construct_transformer(model_args: ModelArgs) -> Transformer:`
`272`	`272`	`norm_eps=model_args.norm_eps,`
`273`	`273`	`)`
`274`	`274`	`)`
	`275`	`+ elif (`
	`276`	`+ model_args.layer_types`
	`277`	`+ and model_args.layer_types[layer_id] == "skip_attention"`
	`278`	`+ ):`
	`279`	`+ attention = AttentionSkip()`
	`280`	`+ transformer_block = TransformerBlock(model_args, attention)`
	`281`	`+ layers.append(transformer_block)`
`275`	`282`	`else:`
`276`	`283`	`attention = cls(`
`277`	`284`	`model_args, layer_id, rope, **model_args.attention_kwargs`