changes for batch size being correct

prishajain1 · prishajain1 · commit a456d10b05fa · 2026-03-02T17:30:39.000+05:30
diff --git a/src/maxdiffusion/configs/ltx2_video.yml b/src/maxdiffusion/configs/ltx2_video.yml
@@ -61,9 +61,9 @@ dcn_data_parallelism: 1  # recommended DCN axis to be auto-sharded
 dcn_fsdp_parallelism: -1
 dcn_context_parallelism: 1
 dcn_tensor_parallelism: 1
-ici_data_parallelism: 1
-ici_fsdp_parallelism: -1  # recommended ICI axis to be auto-sharded
-ici_context_parallelism: 1
+ici_data_parallelism: 2
+ici_fsdp_parallelism: 1  # recommended ICI axis to be auto-sharded
+ici_context_parallelism: 4
 ici_tensor_parallelism: 1
 enable_profiler: False
 
@@ -81,7 +81,7 @@ dataset_name: 'diffusers/pokemon-gpt4-captions'
 train_split: 'train'
 dataset_type: 'tfrecord'
 cache_latents_text_encoder_outputs: True
-per_device_batch_size: 1
+per_device_batch_size: 0.25
 compile_topology_num_slices: -1 
 quantization_local_shard_count: -1
 use_qwix_quantization: False
diff --git a/src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py b/src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py
@@ -1110,8 +1110,7 @@ def __call__(
       )
 
       # 3. Prepare latents
-      _bs = prompt_embeds[0].shape[0] if isinstance(prompt_embeds, list) else prompt_embeds.shape[0]
-      batch_size = _bs // 2 if guidance_scale > 1.0 else _bs
+      batch_size = prompt_embeds[0].shape[0] if isinstance(prompt_embeds, list) else prompt_embeds.shape[0]
       
       # Prepare generators
       if generator is None:

Original file line number	Diff line number	Diff line change
`@@ -1110,8 +1110,7 @@ def __call__(`
`1110`	`1110`	`)`
`1111`	`1111`
`1112`	`1112`	`# 3. Prepare latents`
`1113`		`- _bs = prompt_embeds[0].shape[0] if isinstance(prompt_embeds, list) else prompt_embeds.shape[0]`
`1114`		`- batch_size = _bs // 2 if guidance_scale > 1.0 else _bs`
	`1113`	`+ batch_size = prompt_embeds[0].shape[0] if isinstance(prompt_embeds, list) else prompt_embeds.shape[0]`
`1115`	`1114`
`1116`	`1115`	`# Prepare generators`
`1117`	`1116`	`if generator is None:`