Revert experiment batch size changes

Author: Matrix-Z97

torchtitan/experiments/forge/engine.py

@@ -133,11 +133,6 @@ def __init__(self, config: Config):
             dp_degree, dp_rank = 1, 0
         self.dp_degree, self.dp_rank = dp_degree, dp_rank
 
-        if config.training.global_batch_size < 0:
-            config.training.global_batch_size = (
-                config.training.local_batch_size * dp_degree
-            )
-
         # take control of garbage collection to avoid stragglers
         self.gc_handler = utils.GarbageCollection(
             gc_freq=config.training.gc_freq, debug=config.training.gc_debug
@@ -186,6 +181,10 @@ def __init__(self, config: Config):
 
         # verify batch sizes
         global_batch_size = config.training.global_batch_size
+        if global_batch_size < 0:
+            # This global batch size results in 1 gradient accumulation
+            # step.
+            global_batch_size = config.training.local_batch_size * dp_degree
         assert global_batch_size > 0
         assert (
             global_batch_size % (config.training.local_batch_size * dp_degree) == 0

torchtitan/experiments/graph_trainer/deepseek_v3/parallelize_autoparallel.py

@@ -161,6 +161,9 @@ def parallelize_autoparallel_deepseekv3(
 
     def input_fn():
         global_batch_size = training.global_batch_size
+        if global_batch_size < 0:
+            dp_degree = parallel_dims.dp_replicate * parallel_dims.dp_shard
+            global_batch_size = training.local_batch_size * dp_degree
         tokens = torch.randint(
             0,
             ap_model.model_args.vocab_size,

torchtitan/experiments/graph_trainer/llama3/parallelize_autoparallel.py

@@ -68,6 +68,9 @@ def parallelize_autoparallel_llama(
 
     def input_fn():
         global_batch_size = training.global_batch_size
+        if global_batch_size < 0:
+            dp_degree = parallel_dims.dp_replicate * parallel_dims.dp_shard
+            global_batch_size = training.local_batch_size * dp_degree
         tokens = torch.randint(
             0,
             model.config.vocab_size,

torchtitan/experiments/torchft/trainer.py

@@ -69,11 +69,6 @@ def __init__(self, config: Config):
         # FT addition: adjust dp info via ft_manager
         batch_degree, batch_rank = self.ft_manager.get_dp_info(batch_degree, batch_rank)
 
-        if config.training.global_batch_size < 0:
-            config.training.global_batch_size = (
-                config.training.local_batch_size * batch_degree
-            )
-
         # take control of garbage collection to avoid stragglers
         self.gc_handler = utils.GarbageCollection(
             gc_freq=config.training.gc_freq, debug=config.training.gc_debug
@@ -165,6 +160,10 @@ def __init__(self, config: Config):
 
         # verify batch sizes
         global_batch_size = config.training.global_batch_size
+        if global_batch_size < 0:
+            # This global batch size results in 1 gradient accumulation
+            # step.
+            global_batch_size = config.training.local_batch_size * batch_degree
         assert global_batch_size > 0
         assert (
             global_batch_size % (config.training.local_batch_size * batch_degree) == 0