Fix hardcoded grain_worker_count in HF data processing pipelines to prevent OOM while running Multimodal SFT via pathway

[chiajunglien]

Description

This PR addresses an issue in hf_data_processing.py where the Grain DataLoader's worker count was strictly hardcoded or omitted during instantiation, ignoring the user-provided config.grain_worker_count.

Problem

• In Multimodal SFT (vision_sft_preprocessing_pipeline), worker_count was hardcoded to 1. Because text/image batches can be extremely large, forcing data loading into a multiprocessing background worker quickly exhausted the container's default /dev/shm (Shared Memory), causing a SIGBUS (exit code -7) crash.
• In Text SFT / Pretraining, make_hf_train_iterator and make_hf_eval_iterator failed to pass config.grain_worker_count down to preprocessing_pipeline,
  silently falling back to the default 1 regardless of user configuration.

Solution

• Plumbed config.grain_worker_count directly into the grain.DataLoader inside vision_sft_preprocessing_pipeline.
• Explicitly passed grain_worker_count=config.grain_worker_count to preprocessing_pipeline inside both train and eval iterators.
  This enables users to properly disable multiprocessing on multimodal/heavy data workflows by passing --grain_worker_count=0 in their XPK commands, completely avoiding the /dev/shm limit.

Tests

• Multi-modal SFT run via train_sft_native on ChartQA dataset correctly respects grain_worker_count=0 without crashing.

Command:

xpk workload create-pathways \
  --cluster=mesa-v6e32-eu \
  --workload=multisft-fix \
  --tpu-type=v6e-32 \
  --num-slices=1 \
  --priority=very-high \
  --project=cienet-cmcs \
  --zone=europe-west4-a \
  --skip-validation \
  --docker-image=[gcr.io/tpu-prod-env-multipod/emma-fix-multimodalsft:latest](http://gcr.io/cienet-cmcs/emma-fix-gemma3-0717:latest) \
  --command="set -xue; \
    export TPU_MIN_LOG_LEVEL=0; \
    export TF_CPP_MIN_LOG_LEVEL=0; \
    export TPU_STDERR_LOG_LEVEL=0; \
    export JAX_PLATFORMS='proxy,cpu'; \
    export JAX_BACKEND_TARGET='grpc://127.0.0.1:29000'; \
    export ENABLE_PATHWAYS_PERSISTENCE='1'; \
    python -m maxtext.trainers.post_train.sft.train_sft_native \
      /deps/src/maxtext/configs/post_train/sft-vision-chartqa.yml \
      run_name=post-20260717T013504-ramdisk \
      model_name=gemma3-4b \
      per_device_batch_size=1 \
      max_prefill_predict_length=1024 \
      max_target_length=2048 \
      steps=5 \
      scan_layers=true \
      async_checkpointing=False \
      attention='dot_product' \
      dataset_type=hf \
      hf_train_files='gs://aireenmei-multipod/dataset/hf/chartqa/train-*' \
      hf_path=parquet \
      base_output_directory=gs://us-central1-emmalien-test-83df3ecd-bucket/gemma3-4b-0721/multimodal/sft \
      load_parameters_path=gs://us-central1-emmalien-test-83df3ecd-bucket/post-20260717T013504/0/items \
      dtype=bfloat16 \
      weight_dtype=bfloat16 \
      sharding_tolerance=0.05 \
      checkpoint_storage_use_zarr3=False \
      checkpoint_storage_use_ocdbt=False \
      enable_single_controller=true \
      grain_worker_count=0"

Log: https://paste.googleplex.com/5331889429807104

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[aireenmei]

Could you add a guard in types.py, so it errors out when user set grain_worker_count > 1 for dataset_type=hf?

[chiajunglien]

Could you add a guard in types.py, so it errors out when user set grain_worker_count > 1 for dataset_type=hf?

Added. Thank you!