Repository: instructlab/instructlab

Issue:
Expose `effective_batch_size` parameter per phase for `ilab model train --strategy lab-multiphase`
Required for phase1 vs. phase2 estimated dataset size requirements.


Discussion:

Patch:
diff --git a/src/instructlab/configuration.py b/src/instructlab/configuration.py
index a8328586b9..331c59b0f8 100644
--- a/src/instructlab/configuration.py
+++ b/src/instructlab/configuration.py
@@ -37,6 +37,7 @@
 ILAB_PACKAGE_NAME = "instructlab"
 CONFIG_FILENAME = "config.yaml"
 CONFIG_VERSION = "1.0.0"
+OPTIONAL_POSITIVE_INTEGER = Field(default=None, gt=0)
 
 
 class STORAGE_DIR_NAMES:
@@ -512,11 +513,13 @@ class _train(BaseModel):
     # lab-multiphase training.
     # TODO: could move into its own object.
     # Not strictly necessary for a correct training object.
-    phased_phase1_num_epochs: int | None = None
-    phased_phase1_samples_per_save: int | None = None
+    phased_phase1_num_epochs: int | None = OPTIONAL_POSITIVE_INTEGER
+    phased_phase1_samples_per_save: int | None = OPTIONAL_POSITIVE_INTEGER
+    phased_phase1_effective_batch_size: int | None = OPTIONAL_POSITIVE_INTEGER
 
-    phased_phase2_num_epochs: int | None = None
-    phased_phase2_samples_per_save: int | None = None
+    phased_phase2_num_epochs: int | None = OPTIONAL_POSITIVE_INTEGER
+    phased_phase2_samples_per_save: int | None = OPTIONAL_POSITIVE_INTEGER
+    phased_phase2_effective_batch_size: int | None = OPTIONAL_POSITIVE_INTEGER
 
     phased_mt_bench_judge: str | None = None
 
@@ -667,8 +670,10 @@ def get_default_config() -> Config:
             is_padding_free=False,
             phased_phase1_num_epochs=10,
             phased_phase1_samples_per_save=25000,
+            phased_phase1_effective_batch_size=128,
             phased_phase2_num_epochs=10,
             phased_phase2_samples_per_save=25000,
+            phased_phase2_effective_batch_size=3840,
             phased_mt_bench_judge=DEFAULTS.DEFAULT_JUDGE_MODEL,
         ),
         evaluate=_evaluate(
diff --git a/src/instructlab/model/train.py b/src/instructlab/model/train.py
index 153b647a7d..8b82ff4018 100644
--- a/src/instructlab/model/train.py
+++ b/src/instructlab/model/train.py
@@ -332,15 +332,21 @@ def clickpath_setup(is_dir: bool) -> click.Path:
 @click.option(
     "--phased-phase1-num-epochs",
     cls=clickext.ConfigOption,
-    type=int,
+    type=click.IntRange(min=1),
     help="Number of epochs to run for first phase of end-to-end training.",
 )
 @click.option(
     "--phased-phase1-samples-per-save",
     cls=clickext.ConfigOption,
-    type=int,
+    type=click.IntRange(min=1),
     help="Number of samples to train on between saves for first phase of end-to-end training.",
 )
+@click.option(
+    "--phased-phase1-effective-batch-size",
+    cls=clickext.ConfigOption,
+    type=click.IntRange(min=1),
+    help="Total size of a training batch over all GPUs for Phase 1",
+)
 @click.option(
     "--phased-phase2-data",
     type=clickpath_setup(is_dir=False),
@@ -349,15 +355,21 @@ def clickpath_setup(is_dir: bool) -> click.Path:
 @click.option(
     "--phased-phase2-num-epochs",
     cls=clickext.ConfigOption,
-    type=int,
+    type=click.IntRange(min=1),
     help="Number of epochs to run for second phase of end-to-end training.",
 )
 @click.option(
     "--phased-phase2-samples-per-save",
     cls=clickext.ConfigOption,
-    type=int,
+    type=click.IntRange(min=1),
     help="Number of samples to train on between saves for second phase of end-to-end training.",
 )
+@click.option(
+    "--phased-phase2-effective-batch-size",
+    cls=clickext.ConfigOption,
+    type=click.IntRange(min=1),
+    help="Total size of a training batch over all GPUs for Phase 2",
+)
 @click.option(
     "--phased-mt-bench-judge",
     # type=clickpath_setup(is_dir=True), # want this in the future, can't guarantee it exists so can't enforce it this way.
@@ -398,9 +410,11 @@ def train(
     phased_phase1_data: pathlib.Path | None,
     phased_phase1_num_epochs: int | None,
     phased_phase1_samples_per_save: int | None,
+    phased_phase1_effective_batch_size: int | None,
     phased_phase2_data: pathlib.Path | None,
     phased_phase2_num_epochs: int | None,
     phased_phase2_samples_per_save: int | None,
+    phased_phase2_effective_batch_size: int | None,
     phased_mt_bench_judge: pathlib.Path | None,
     skip_user_confirm: bool,
     enable_serving_output: bool,
@@ -671,10 +685,12 @@ def get_files(directory: str, pattern: str) -> list[str]:
                 phase1_num_epochs=phased_phase1_num_epochs,
                 phase1_samples_per_save=phased_phase1_samples_per_save,
                 phase1_checkpoints_dir=phased_base_dir / "phase1" / "checkpoints",
+                phased_phase1_effective_batch_size=phased_phase1_effective_batch_size,
                 phase2_data=phased_phase2_data,
                 phase2_num_epochs=phased_phase2_num_epochs,
                 phase2_samples_per_save=phased_phase2_samples_per_save,
                 phase2_checkpoints_dir=phased_base_dir / "phase2" / "checkpoints",
+                phased_phase2_effective_batch_size=phased_phase2_effective_batch_size,
                 phase2_eval_cache=phased_base_dir / "phase2" / "eval_cache",
                 mtbench_judge=mt_bench_judge,
                 enable_serving_output=enable_serving_output,
@@ -729,6 +745,7 @@ def _training_phase(
     num_epochs: int | None = None,
     samples_per_save: int | None = None,
     checkpoint_dir: pathlib.Path | None = None,
+    effective_batch_size: int | None = None,
 ) -> None:
     """A single step of phased training that supports key param overriding."""
 
@@ -763,6 +780,12 @@ def _training_phase(
         )
         train_args.save_samples = samples_per_save
 
+    if effective_batch_size:
+        logger.debug(
+            f"Phased Training -- training phase -- Overriding effective batch size: {train_args.effective_batch_size} with {effective_batch_size}"
+        )
+        train_args.effective_batch_size = effective_batch_size
+
     click.secho(
         f"TrainingArgs for current phase: {pprint.pformat(train_args)}", fg="cyan"
     )
@@ -900,10 +923,12 @@ def _run_phased_training(
     phase1_num_epochs: int | None,
     phase1_samples_per_save: int | None,
     phase1_checkpoints_dir: pathlib.Path,
+    phased_phase1_effective_batch_size: int | None,
     phase2_data: pathlib.Path,
     phase2_num_epochs: int | None,
     phase2_samples_per_save: int | None,
     phase2_checkpoints_dir: pathlib.Path,
+    phased_phase2_effective_batch_size: int | None,
     phase2_eval_cache: pathlib.Path,
     mtbench_judge: pathlib.Path,
     enable_serving_output: bool,
@@ -917,6 +942,7 @@ def _run_phased_training(
         checkpoint_dir=phase1_checkpoints_dir,
         num_epochs=phase1_num_epochs,
         samples_per_save=phase1_samples_per_save,
+        effective_batch_size=phased_phase1_effective_batch_size,
         # model override not necessary because we expect model to come from ctx.params.model_path.
     )
 
@@ -936,6 +962,7 @@ def _run_phased_training(
         model_override=p1_best_ckpt,
         num_epochs=phase2_num_epochs,
         samples_per_save=phase2_samples_per_save,
+        effective_batch_size=phased_phase2_effective_batch_size,
     )
 
     click.secho("MT-Bench evaluation for Phase 2...", fg="cyan")


Test patch:
diff --git a/tests/testdata/default_config.yaml b/tests/testdata/default_config.yaml
index 30a7d4503d..b05d817e99 100644
--- a/tests/testdata/default_config.yaml
+++ b/tests/testdata/default_config.yaml
@@ -48,9 +48,9 @@ generate:
       max_ctx_size: 4096
     model_path: /cache/instructlab/models/merlinite-7b-lab-Q4_K_M.gguf
     vllm:
+      gpus: null
       llm_family: ''
       max_startup_attempts: 300
-      gpus: null
       vllm_args: []
 serve:
   backend: null
@@ -62,9 +62,9 @@ serve:
     max_ctx_size: 4096
   model_path: /cache/instructlab/models/merlinite-7b-lab-Q4_K_M.gguf
   vllm:
+    gpus: null
     llm_family: ''
     max_startup_attempts: 300
-    gpus: null
     vllm_args: []
 train:
   additional_args: {}
@@ -82,8 +82,10 @@ train:
   nproc_per_node: 1
   num_epochs: 10
   phased_mt_bench_judge: /cache/instructlab/models/prometheus-eval/prometheus-8x7b-v2.0
+  phased_phase1_effective_batch_size: 128
   phased_phase1_num_epochs: 10
   phased_phase1_samples_per_save: 25000
+  phased_phase2_effective_batch_size: 3840
   phased_phase2_num_epochs: 10
   phased_phase2_samples_per_save: 25000
   save_samples: 250000