Train & Eval on entire dataset when pdbs < 1

[muskansh-google]

Description

This PR fixes a bug where data was silently dropped when running with fractional per-device batch sizes (per_device_batch_size < 1.0 or eval_per_device_batch_size < 1.0). In the past, the dataloader loaded at least one sample per device, but loss_fn truncated the input batch down to micro_batch_size_to_train_on or micro_batch_size_to_eval_on, discarding the remaining samples. With this change, the slicing logic is removed, and all loaded samples are processed across microbatches using gradient accumulation during training and forward-only metric accumulation during evaluation.

Why This Solution

Instead of discarding data, setting per_device_batch_size < 1.0 is treated as implicit microbatch accumulation:

• For training, gradient accumulation is automatically computed over num_microbatches = global_batch_size_to_load // micro_batch_size_to_train_on. All loaded samples are processed across microbatches and accumulated into the gradients.
• For evaluation, a dedicated forward-only scan loop (_fractional_batch_eval) processes each microbatch and aggregates metrics without performing unnecessary backpropagation.
• No data is thrown away, batch size calculations are consistent, and throughput matches the loaded dataset volume.

Tests

Added unit tests to verify batch sizing calculations and microbatch accumulation:

• tests/unit/pyconfig_test.py:
  • test_calculate_global_batch_sizes_unified: Verifies calculation logic across fractional batch sizes, gradient accumulation steps, and expansion factors.
  • test_fractional_per_device_batch_size_calculations: Verifies config initialization with fractional train and eval batch sizes.
• tests/unit/gradient_accumulation_nnx_test.py:
  • test_fractional_per_device_batch_size_accumulates_all_samples: Verifies gradient accumulation across all 4 samples without data dropping when per_device_batch_size = 0.25.
  • test_fractional_eval_microbatch_reshaping_and_accumulation: Verifies microbatch reshaping and metric accumulation for eval without data dropping.

Commands to run the tests:

pytest tests/unit/pyconfig_test.py -k "calculate_global_batch_sizes_unified or fractional_per_device_batch_size"
pytest tests/unit/gradient_accumulation_nnx_test.py -k "fractional"

[gemini-code-assist]

Code Review

This pull request refactors batch size calculations and execution paths to properly support fractional batch sizes (where per-device batch size is less than 1.0) during both training and evaluation. Instead of decimating data, the code now leverages gradient accumulation or a new fractional evaluation helper to accumulate metrics across micro-batches. The review feedback suggests centralizing the duplicated batch size calculation helper function to a shared utility module and refactoring the hardcoded auxiliary loss keys in the fractional evaluation helper into a list constant to improve maintainability.

[gemini-code-assist]

This calculate_global_batch_sizes helper function duplicates the logic from pyconfig_deprecated.py:calculate_global_batch_sizes. To improve maintainability and prevent future inconsistencies, consider centralizing this logic.

You could move the standalone calculate_global_batch_sizes function from pyconfig_deprecated.py to a shared utility module (e.g., in maxtext/utils/) and import it here and in pyconfig_deprecated.py.

[gemini-code-assist]

The init_acc dictionary and the accumulate_eval function hardcode the names of auxiliary losses to be accumulated (e.g., z_loss, moe_lb_loss). This makes the function brittle; if a new auxiliary loss is added to loss_fn, it will be silently ignored during fractional evaluation unless this function is also updated.

To improve maintainability, you could define the list of keys to accumulate in a single place. This makes it easier to add new losses in the future.

For example:

# At the top of the file or as a local constant
ACCUMULATED_EVAL_KEYS = [
    "xent_sum",
    "total_weights",
    "z_loss",
    "moe_lb_loss",
    "indexer_loss",
    "mtp_loss",
]

# ... inside _fractional_batch_eval ...
def accumulate_eval(acc, micro_batch):
  _, aux = single_eval_fn(micro_batch)
  new_acc = {k: acc[k] + aux.get(k, 0.0) for k in ACCUMULATED_EVAL_KEYS}
  return new_acc, None

init_acc = {k: 0.0 for k in ACCUMULATED_EVAL_KEYS}

This would make the accumulation logic more robust to future changes.