sharding keys across context

prishajain1 · prishajain1 · commit 66630fbda40b · 2026-04-16T19:35:46.000+05:30
diff --git a/src/maxdiffusion/configs/ltx2_video.yml b/src/maxdiffusion/configs/ltx2_video.yml
@@ -55,6 +55,7 @@ logical_axis_rules: [
                       ['conv_batch', ['data', 'context', 'fsdp']],
                       ['out_channels', 'tensor'],
                       ['conv_out', 'context'],
+                      ['activation_kv_length', 'context'],
                     ]
 data_sharding: ['data', 'fsdp', 'context', 'tensor']
 

Original file line number	Diff line number	Diff line change
`@@ -55,6 +55,7 @@ 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`	`59`	`]`
`59`	`60`	`data_sharding: ['data', 'fsdp', 'context', 'tensor']`
`60`	`61`