sharding keys across context

prishajain1 · prishajain1 · commit bb658348b614 · 2026-04-16T20:13:37.000+05:30
diff --git a/src/maxdiffusion/configs/ltx2_video.yml b/src/maxdiffusion/configs/ltx2_video.yml
@@ -55,7 +55,8 @@ logical_axis_rules: [
                       ['conv_batch', ['data', 'context', 'fsdp']],
                       ['out_channels', 'tensor'],
                       ['conv_out', 'context'],
-                      ['activation_kv_length', 'context'],
+                      ['activation_self_attn_q_length', 'context'],
+                      ['activation_self_attn_kv_length', 'context'],
                     ]
 data_sharding: ['data', 'fsdp', 'context', 'tensor']
 

Original file line number	Diff line number	Diff line change
`@@ -55,7 +55,8 @@ logical_axis_rules: [`
`55`	`55`	`['conv_batch', ['data', 'context', 'fsdp']],`
`56`	`56`	`['out_channels', 'tensor'],`
`57`	`57`	`['conv_out', 'context'],`
`58`		`- ['activation_kv_length', 'context'],`
	`58`	`+ ['activation_self_attn_q_length', 'context'],`
	`59`	`+ ['activation_self_attn_kv_length', 'context'],`
`59`	`60`	`]`
`60`	`61`	`data_sharding: ['data', 'fsdp', 'context', 'tensor']`
`61`	`62`