Merge pull request #3758 from AI-Hypercomputer:aireen/mixture_expansion

Google-ML-Automation · Google-ML-Automation · commit a4bd69d7b6d9 · 2026-04-29T11:43:01.000-07:00
PiperOrigin-RevId: 907699767
diff --git a/src/maxtext/input_pipeline/grain_data_processing.py b/src/maxtext/input_pipeline/grain_data_processing.py
@@ -424,6 +424,7 @@ def make_grain_train_iterator(
         grain_num_threads=config.grain_num_threads,
         grain_prefetch_buffer_size=config.grain_prefetch_buffer_size,
         grain_data_source_max_workers=config.grain_data_source_max_workers,
+        mixture_config_path=config.grain_train_mixture_config_path,
     )
     if config.use_dpo:
       preprocessing_fn = functools.partial(

Original file line number	Diff line number	Diff line change
`@@ -424,6 +424,7 @@ def make_grain_train_iterator(`
`424`	`424`	`grain_num_threads=config.grain_num_threads,`
`425`	`425`	`grain_prefetch_buffer_size=config.grain_prefetch_buffer_size,`
`426`	`426`	`grain_data_source_max_workers=config.grain_data_source_max_workers,`
	`427`	`+ mixture_config_path=config.grain_train_mixture_config_path,`
`427`	`428`	`)`
`428`	`429`	`if config.use_dpo:`
`429`	`430`	`preprocessing_fn = functools.partial(`