AI-Hypercomputer / AI-Hypercomputer/maxtext

Rampup batch size settings are not validated in the Pydantic config

未关闭
#4,690 0 条评论 0 个 reaction 已指派 1 人 已被 @NuojCheng 认领 在 GitHub 查看
主要语言
Python
星标
2.4k
派生
607
平均合并
2 天 19 小时
30 天内合并 PR
158

描述

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.

贡献指南

打开贡献指南

评估

这个 Issue 还没有评估数据。

把新 issue 发到你的邮箱

精选适合新手参与的 GitHub issue 摘要。