The legacy config validates the rampup batch size settings, the Pydantic config does not. Invalid combinations are accepted and either switch rampup off with no diagnostic, or ramp to a batch size that was never requested.
pyconfig_deprecated.validate_rampup_batch_size (src/maxtext/configs/pyconfig_deprecated.py:203) asserts five things: per_device_batch_size_start > 0, per_device_batch_size_increment > 0, global_rampup_samples > 0, per_device_batch_size - per_device_batch_size_start > 0, and that the difference divides by the increment.
DatasetGeneral in src/maxtext/configs/types.py declares the same four fields around line 1365 with no validator. The schedule derivation in MaxTextConfig around line 3067 is written defensively, if self.global_batch_size_to_load_increment > 0: then if num_increments > 0:, so when either guard fails the block falls through and leaves rampup_end_step = 0.
Replicating that arithmetic for 8 devices, expansion_factor_real_data=1, gradient_accumulation_steps=1:
| settings |
rampup_end_step |
result |
per_device_batch_size=8, start=4, increment=2, samples=500 |
14 |
correct |
increment=0 |
0 |
rampup silently off |
global_rampup_samples=0 |
0 |
rampup silently off |
start=8, per_device_batch_size=4 |
0 |
rampup silently off |
per_device_batch_size=9, start=4, increment=2 |
14 |
ramps to global batch size 64, 72 was requested |
The first three mean a run configured with enable_rampup_batch_size=True trains at a constant batch size with nothing in the logs saying rampup never engaged. The last one is worse, num_increments = diff // increment truncates, so ramp-up finishes below the configured per_device_batch_size and stays there.
The Pydantic config should reject these at parse time the way the legacy path does. I can send a PR adding a model_validator(mode="after") on DatasetGeneral with unit tests in tests/unit/configs_value_test.py.
The legacy config validates the rampup batch size settings, the Pydantic config does not. Invalid combinations are accepted and either switch rampup off with no diagnostic, or ramp to a batch size that was never requested.
pyconfig_deprecated.validate_rampup_batch_size(src/maxtext/configs/pyconfig_deprecated.py:203) asserts five things:per_device_batch_size_start > 0,per_device_batch_size_increment > 0,global_rampup_samples > 0,per_device_batch_size - per_device_batch_size_start > 0, and that the difference divides by the increment.DatasetGeneralin src/maxtext/configs/types.py declares the same four fields around line 1365 with no validator. The schedule derivation inMaxTextConfigaround line 3067 is written defensively,if self.global_batch_size_to_load_increment > 0:thenif num_increments > 0:, so when either guard fails the block falls through and leavesrampup_end_step = 0.Replicating that arithmetic for 8 devices,
expansion_factor_real_data=1,gradient_accumulation_steps=1:per_device_batch_size=8, start=4, increment=2, samples=500increment=0global_rampup_samples=0start=8, per_device_batch_size=4per_device_batch_size=9, start=4, increment=2The first three mean a run configured with
enable_rampup_batch_size=Truetrains at a constant batch size with nothing in the logs saying rampup never engaged. The last one is worse,num_increments = diff // incrementtruncates, so ramp-up finishes below the configuredper_device_batch_sizeand stays there.The Pydantic config should reject these at parse time the way the legacy path does. I can send a PR adding a
model_validator(mode="after")onDatasetGeneralwith unit tests in tests/unit/configs_value_test.py.