Skip to content

Validate rampup batch size settings in the Pydantic config - #4691

Open
Atishyy27 wants to merge 1 commit into
AI-Hypercomputer:mainfrom
Atishyy27:fix/rampup-batch-size-validation
Open

Validate rampup batch size settings in the Pydantic config#4691
Atishyy27 wants to merge 1 commit into
AI-Hypercomputer:mainfrom
Atishyy27:fix/rampup-batch-size-validation

Conversation

@Atishyy27

Copy link
Copy Markdown

Closes #4690.

The legacy config validates the rampup batch size settings in pyconfig_deprecated.validate_rampup_batch_size (src/maxtext/configs/pyconfig_deprecated.py:203). The Pydantic config has no equivalent, so an unusable combination is accepted and then silently degrades.

The derivation in MaxTextConfig is guarded by if self.global_batch_size_to_load_increment > 0: and if num_increments > 0:, so a non-positive increment, non-positive global_rampup_samples, or a start at or above per_device_batch_size falls through and leaves rampup_end_step = 0. Rampup is off and nothing says so. When the change is not a multiple of the increment, num_increments = diff // increment truncates and ramp-up finishes below the configured per_device_batch_size.

Replicating the derivation 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 requested

This adds a model_validator(mode="after") on DatasetGeneral with the same five checks as the legacy path, skipped entirely when enable_rampup_batch_size is False.

Tests in tests/unit/configs_value_test.py cover each rejection, a positive control asserting a well-formed config still produces rampup_end_step > 0, and a case confirming the checks do not fire when rampup is disabled.

The legacy config path validates these in
pyconfig_deprecated.validate_rampup_batch_size, but the Pydantic config
has no equivalent, so an unusable combination is accepted silently.

The schedule derivation in MaxTextConfig is guarded by
'if global_batch_size_to_load_increment > 0' and 'if num_increments > 0',
so a non-positive increment, non-positive global_rampup_samples, or a
start at or above per_device_batch_size falls through and leaves
rampup_end_step = 0, disabling rampup with no diagnostic. When the batch
size change is not a multiple of the increment, the floor division
truncates and ramp-up finishes at a smaller batch size than configured.

Add a model_validator on DatasetGeneral mirroring the legacy checks, and
cover each rejection plus the disabled-rampup case in
tests/unit/configs_value_test.py.
@gemini-code-assist

Copy link
Copy Markdown

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Rampup batch size settings are not validated in the Pydantic config

1 participant