Merge pull request #139 from baidu-baige/feature/vlm-data-shuffling

[vlm, data] feat: expose Energon read-order controls

Author: kaimo455

loongforge/data/multimodal/dataloader_provider.py

@@ -21,7 +21,7 @@
 from megatron.core.transformer.enums import AttnMaskType
 from megatron.training import get_args
 from megatron.training.checkpointing import get_checkpoint_name
-from loongforge.utils import constants, get_model_config
+from loongforge.utils import constants, get_model_config, print_rank_0
 from .base.task_encoder import print_error_handler
 from loongforge.train.get_position_idx_func import get_position_ids
 
@@ -307,6 +307,26 @@ def _build_masks_and_positions(self, batch: Dict[str, Any]) -> None:
         batch["loss_mask"] = loss_mask
 
 
+def _energon_read_order_kwargs(args):
+    """Resolve optional Energon sample-order controls from CLI flags.
+
+    ``shuffle_buffer_size`` adds sample-level shuffling before cooking and
+    encoding. For an offline-packed dataset, each raw sample is a complete pack,
+    so this randomizes pack order without changing pack contents.
+    ``max_samples_per_sequence`` targets shorter runs from each shard slice,
+    giving Energon more sequences to mix through its parallel shard iterators.
+    Disabled values are passed as ``None`` to preserve Energon's defaults.
+    """
+    shuffle_buffer_size = getattr(args, "data_shuffle_buffer_size", 0) or 0
+    max_samples_per_sequence = getattr(args, "data_max_samples_per_sequence", 0) or 0
+    return {
+        "shuffle_buffer_size": shuffle_buffer_size if shuffle_buffer_size > 1 else None,
+        "max_samples_per_sequence": (
+            max_samples_per_sequence if max_samples_per_sequence > 0 else None
+        ),
+    }
+
+
 def get_train_dataset(task_encoder):
     """Get the training dataset"""
     args = get_args()
@@ -318,18 +338,24 @@ def get_train_dataset(task_encoder):
         worker_debug_path=None,
         worker_log_level=0,
     )
+    shuffle_kwargs = _energon_read_order_kwargs(args)
+    print_rank_0(
+        f"> Energon read order: shuffle_buffer_size="
+        f"{shuffle_kwargs['shuffle_buffer_size']}, max_samples_per_sequence="
+        f"{shuffle_kwargs['max_samples_per_sequence']}",
+        args.rank,
+    )
 
     if len(args.data_path) == 1:
         train_ds = energon.get_train_dataset(
             args.data_path[0],
             batch_size=args.micro_batch_size,
             task_encoder=task_encoder,
             worker_config=worker_config,
-            max_samples_per_sequence=None,
-            shuffle_buffer_size=None,
             packing_buffer_size=args.packing_buffer_size,
             handler=print_error_handler,
             image_decode="pil",
+            **shuffle_kwargs,
         )
     else:
         data_paths, data_weights = get_blend_from_list(args.data_path)
@@ -339,11 +365,10 @@ def get_train_dataset(task_encoder):
             batch_size=args.micro_batch_size,
             task_encoder=task_encoder,
             worker_config=worker_config,
-            max_samples_per_sequence=None,
-            shuffle_buffer_size=None,
             packing_buffer_size=args.packing_buffer_size,
             handler=print_error_handler,
             image_decode="pil",
+            **shuffle_kwargs,
         )
     return train_ds
 

loongforge/train/arguments.py

@@ -1182,6 +1182,29 @@ def _add_extra_multimodal_args(parser):
         help="Path to save Energon dataloader state for resumable training. Default: None"
     )
 
+    group.add_argument(
+        "--data-shuffle-buffer-size",
+        type=int,
+        default=0,
+        help="Size of the Energon sample shuffle buffer applied before cooking and "
+             "encoding. For an offline-packed dataset, each sample is one complete "
+             "packed sequence, so this randomizes pack order without changing pack "
+             "contents. Use this to reduce length correlation when shards are ordered "
+             "by sample length. Larger buffers increase cross-sample mixing and worker "
+             "memory use. 0 or 1 disables the sample-level buffer (default)."
+    )
+
+    group.add_argument(
+        "--data-max-samples-per-sequence",
+        type=int,
+        default=0,
+        help="Target size of each shard sequence used by Energon for parallel "
+             "iteration. Large worker shard ranges are divided into sequences of roughly "
+             "this many samples, creating more sequences that Energon can interleave. "
+             "Use this when --data-shuffle-buffer-size cannot cover long ordered runs. "
+             "0 leaves Energon's limit unset (default)."
+    )
+
     group.add_argument(
         "--packing-pretrain-data",
         action="store_true",

tests/test_vlm_dataloader_read_order.py

@@ -0,0 +1,148 @@
+# Copyright 2026 The LoongForge Authors.
+# SPDX-License-Identifier: Apache-2.0
+
+import argparse
+import ast
+from pathlib import Path
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import pytest
+
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+
+
+def _load_functions(relative_path, names, namespace=None):
+    """Load selected pure-Python functions without importing the GPU stack."""
+    path = REPO_ROOT / relative_path
+    tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path))
+    functions = [
+        node
+        for node in tree.body
+        if isinstance(node, ast.FunctionDef) and node.name in names
+    ]
+    assert {function.name for function in functions} == set(names)
+
+    namespace = {} if namespace is None else namespace
+    module = ast.Module(body=functions, type_ignores=[])
+    exec(compile(module, str(path), "exec"), namespace)
+    return namespace
+
+
+@pytest.mark.parametrize(
+    ("args", "expected"),
+    [
+        (
+            SimpleNamespace(),
+            {"shuffle_buffer_size": None, "max_samples_per_sequence": None},
+        ),
+        (
+            SimpleNamespace(
+                data_shuffle_buffer_size=1,
+                data_max_samples_per_sequence=0,
+            ),
+            {"shuffle_buffer_size": None, "max_samples_per_sequence": None},
+        ),
+        (
+            SimpleNamespace(
+                data_shuffle_buffer_size=4096,
+                data_max_samples_per_sequence=256,
+            ),
+            {"shuffle_buffer_size": 4096, "max_samples_per_sequence": 256},
+        ),
+    ],
+)
+def test_energon_read_order_kwargs(args, expected):
+    namespace = _load_functions(
+        "loongforge/data/multimodal/dataloader_provider.py",
+        {"_energon_read_order_kwargs"},
+    )
+    assert namespace["_energon_read_order_kwargs"](args) == expected
+
+
+@pytest.mark.parametrize(
+    ("data_path", "expected_path"),
+    [
+        (["dataset"], "dataset"),
+        (["dataset-a", "dataset-b"], "/tmp/metadataset.yaml"),
+    ],
+)
+def test_get_train_dataset_forwards_read_order_kwargs(data_path, expected_path):
+    get_train_dataset = Mock(return_value="train-dataset")
+    worker_config = object()
+    energon = SimpleNamespace(
+        WorkerConfig=Mock(return_value=worker_config),
+        get_train_dataset=get_train_dataset,
+    )
+    args = SimpleNamespace(
+        data_path=data_path,
+        micro_batch_size=2,
+        num_workers=4,
+        packing_buffer_size=10000,
+        data_shuffle_buffer_size=4096,
+        data_max_samples_per_sequence=256,
+        rank=0,
+    )
+    namespace = {
+        "energon": energon,
+        "parallel_state": SimpleNamespace(
+            get_data_parallel_rank=lambda: 2,
+            get_data_parallel_world_size=lambda: 8,
+            get_data_parallel_group=lambda: "dp-group",
+        ),
+        "get_args": lambda: args,
+        "get_blend_from_list": Mock(
+            return_value=(["dataset-a", "dataset-b"], [0.25, 0.75])
+        ),
+        "create_metadataset_yaml": Mock(return_value="/tmp/metadataset.yaml"),
+        "print_error_handler": object(),
+        "print_rank_0": Mock(),
+    }
+    namespace = _load_functions(
+        "loongforge/data/multimodal/dataloader_provider.py",
+        {"_energon_read_order_kwargs", "get_train_dataset"},
+        namespace,
+    )
+
+    assert namespace["get_train_dataset"]("task-encoder") == "train-dataset"
+    get_train_dataset.assert_called_once()
+    path, = get_train_dataset.call_args.args
+    kwargs = get_train_dataset.call_args.kwargs
+    assert path == expected_path
+    assert kwargs["shuffle_buffer_size"] == 4096
+    assert kwargs["max_samples_per_sequence"] == 256
+    assert kwargs["task_encoder"] == "task-encoder"
+    assert kwargs["worker_config"] is worker_config
+
+
+def test_multimodal_argument_defaults_and_overrides():
+    class LanguageModelFamilies:
+        @staticmethod
+        def names():
+            return ["llama"]
+
+    namespace = _load_functions(
+        "loongforge/train/arguments.py",
+        {"_add_extra_multimodal_args"},
+        {
+            "get_support_model_archs": lambda values: values,
+            "constants": SimpleNamespace(LanguageModelFamilies=LanguageModelFamilies),
+        },
+    )
+    parser = namespace["_add_extra_multimodal_args"](argparse.ArgumentParser())
+
+    defaults = parser.parse_args([])
+    assert defaults.data_shuffle_buffer_size == 0
+    assert defaults.data_max_samples_per_sequence == 0
+
+    configured = parser.parse_args(
+        [
+            "--data-shuffle-buffer-size",
+            "4096",
+            "--data-max-samples-per-sequence",
+            "256",
+        ]
+    )
+    assert configured.data_shuffle_buffer_size == 4096
+    assert configured.data_max_samples_per_sequence == 256