fix(data): address Energon tutorial review

Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>

Author: yaoyu-33

tutorials/data/energon/nemotron-image-v3.md

@@ -101,41 +101,36 @@ export PRETRAINED_CHECKPOINT=/checkpoints/Qwen3-VL-8B-Instruct
 
 ## 4. Launch Qwen3-VL LoRA through `train.sh`
 
-The unified launcher selects the Qwen3-VL Energon task encoder with `--dataset energon`; the exact Energon recipe
-selects `vlm_step`. Mount the prepared dataset, checkpoint, and output paths into the training container:
+The exact recipe already selects LoRA, the Energon dataset config, and `vlm_step`; those launcher selectors do not need
+to be repeated. Mount the prepared dataset, checkpoint, and output paths into the training container:
 
 ```bash
 export OUTPUT_DIR=/results/qwen3-vl-nemotron-image-v3
 mkdir -p "$OUTPUT_DIR"
 
 ./scripts/training/train.sh \
-  --nodes 1 --gpus-per-node 1 \
+  --gpus-per-node 1 \
   --account ACCOUNT --partition PARTITION \
   --container-image /path/to/container.sqsh \
   --mount "$ENERGON_PATH" \
   --mount "$PRETRAINED_CHECKPOINT" \
   --mount "$OUTPUT_DIR" \
   --recipe qwen3_vl_8b_peft_energon_config \
-  --step-func vlm_step \
-  --mode lora --dataset energon \
   --seq_length 16384 \
-  checkpoint.pretrained_checkpoint="$PRETRAINED_CHECKPOINT" \
+  --pretrained_checkpoint "$PRETRAINED_CHECKPOINT" \
+  --save_dir "$OUTPUT_DIR/checkpoints" \
+  --save_interval 1 \
+  --max_steps 1 \
+  --global_batch_size 1 \
+  --micro_batch_size 1 \
   checkpoint.load=null \
-  checkpoint.save="$OUTPUT_DIR/checkpoints" \
-  checkpoint.save_interval=1 \
-  train.train_iters=1 \
-  train.global_batch_size=1 \
-  train.micro_batch_size=1 \
   validation.eval_interval=1 \
   validation.eval_iters=1 \
-  validation.eval_micro_batch_size=1 \
   dataset.path="$ENERGON_PATH" \
   dataset.micro_batch_size=1 \
   dataset.num_workers=0 \
-  dataset.num_val_workers=0 \
   dataset.packing_buffer_size=16 \
   dataset.task_encoder.hf_processor_revision="$MODEL_REVISION" \
-  +tokenizer.hf_tokenizer_kwargs.revision="$MODEL_REVISION" \
   model.recompute_granularity=full \
   model.recompute_method=uniform \
   model.recompute_num_layers=1 \
@@ -148,10 +143,16 @@ mkdir -p "$OUTPUT_DIR"
 The pinned `turing` subset contains long reasoning targets: with the pinned Qwen3-VL processor and default image
 pixel bounds, the longest observed row is 15,608 tokens. The explicit 16,384-token setting keeps every source row
 eligible; the recipe's 4,096-token default would skip overlength rows. Full activation recompute keeps the documented
-one-GPU LoRA smoke within an 80 GiB H100. This one-step baseline intentionally uses no worker subprocesses so its
-command matches the validated smoke exactly. After it passes, increase `train.train_iters`, global batch size, and
-`dataset.num_workers`/`dataset.num_val_workers` for the target run. Benchmark recompute against context parallelism on
-the target hardware instead of silently reducing the dataset.
+one-GPU LoRA smoke within an 80 GiB H100. Native packing requires physical micro-batch size 1 in both the training and
+Energon configs, per-token loss, and non-averaged DDP collectives; these settings are not optional for this packed run.
+The processor revision is also required to reproduce the measured sequence lengths and pins the processor's tokenizer
+artifacts to the same model revision.
+
+The one-step limit, global batch size 1, validation/checkpoint interval 1, worker count 0, and log interval 1 are
+smoke/debug settings rather than packing requirements. They preserve the validated baseline and make its loss and
+checkpoint evidence visible. After it passes, increase the step count, global batch size, and worker count for the
+target run. `num_val_workers` defaults to `num_workers`. Benchmark recompute against context parallelism on the target
+hardware instead of silently reducing the dataset.
 
 For a local interactive smoke, export the same `ENERGON_PATH` and `PRETRAINED_CHECKPOINT`, then run
 `bash examples/models/qwen/qwen3_vl/peft_energon.sh`.

tutorials/data/energon/prepare_nemotron_image_v3.py

@@ -25,7 +25,6 @@
 
 import argparse
 import hashlib
-import io
 import json
 import logging
 import pickle
@@ -36,6 +35,8 @@
 from typing import Any
 from urllib.parse import urlparse
 
+import webdataset as wds
+
 from megatron.bridge.data.energon import prepare_webdataset
 
 
@@ -197,52 +198,34 @@ def read(self, reference: str) -> bytes:
         )
 
 
-class SplitShardWriter:
-    """Write deterministic split-prefixed WebDataset tar shards."""
-
-    def __init__(self, output_dir: Path, split: str, max_samples_per_tar: int) -> None:
-        self.output_dir = output_dir
-        self.split = split
-        self.max_samples_per_tar = max_samples_per_tar
-        self._archive: tarfile.TarFile | None = None
-        self._shard_index = 0
-        self._samples_in_shard = 0
-
-    def close(self) -> None:
-        """Close the current shard, if one was opened."""
-        if self._archive is not None:
-            self._archive.close()
-            self._archive = None
+def _make_shard_writer(output_dir: Path, split: str, max_samples_per_tar: int) -> wds.ShardWriter:
+    """Create the standard WebDataset writer with deterministic tar metadata."""
+    return wds.ShardWriter(
+        str(output_dir / f"{split}-shard-%06d.tar"),
+        maxcount=max_samples_per_tar,
+        maxsize=float("inf"),
+        verbose=0,
+        encoder=False,
+        mtime=0,
+        mode=0o644,
+        user="",
+        group="",
+    )
 
-    def __enter__(self) -> SplitShardWriter:
-        return self
 
-    def __exit__(self, *args: object) -> None:
-        self.close()
-
-    def _open_shard(self) -> None:
-        self.close()
-        shard_path = self.output_dir / f"{self.split}-shard-{self._shard_index:06d}.tar"
-        self._archive = tarfile.open(shard_path, "w")
-        self._shard_index += 1
-        self._samples_in_shard = 0
-
-    def _add_bytes(self, name: str, payload: bytes) -> None:
-        if self._archive is None:
-            raise RuntimeError("Shard writer has no open archive.")
-        info = tarfile.TarInfo(name=name)
-        info.size = len(payload)
-        info.mtime = 0
-        info.mode = 0o644
-        self._archive.addfile(info, io.BytesIO(payload))
-
-    def write(self, key: str, images: list[bytes], messages: list[dict[str, Any]]) -> None:
-        """Write one matching-key media list and conversation."""
-        if self._archive is None or self._samples_in_shard >= self.max_samples_per_tar:
-            self._open_shard()
-        self._add_bytes(f"{key}.jpgs", pickle.dumps(images, protocol=4))
-        self._add_bytes(f"{key}.json", json.dumps(messages, ensure_ascii=False).encode("utf-8"))
-        self._samples_in_shard += 1
+def _write_sample(
+    writer: wds.ShardWriter,
+    key: str,
+    images: list[bytes],
+    messages: list[dict[str, Any]],
+) -> None:
+    writer.write(
+        {
+            "__key__": key,
+            "jpgs": pickle.dumps(images, protocol=4),
+            "json": json.dumps(messages, ensure_ascii=False).encode("utf-8"),
+        }
+    )
 
 
 def _sample_split(key: str, validation_fraction: float) -> str:
@@ -357,46 +340,49 @@ def convert_subsets(
     seen_keys: set[str] = set()
     total = 0
 
-    with SplitShardWriter(output_dir, "train", max_samples_per_tar) as train_writer:
-        with SplitShardWriter(output_dir, "val", max_samples_per_tar) as val_writer:
-            writers = {"train": train_writer, "val": val_writer}
-            for subset in subsets:
-                subset_dir = (source_dir / subset).resolve()
-                if not subset_dir.is_relative_to(source_dir) or not subset_dir.is_dir():
-                    raise FileNotFoundError(f"Downloaded subset directory does not exist: {subset_dir}")
-                jsonl_path = _jsonl_path(subset_dir, subset)
-                with (
-                    MediaResolver(
-                        subset_dir,
-                        allow_loose_media=allow_loose_media,
-                        archive_paths=None if media_archives is None else media_archives.get(subset, ()),
-                    ) as resolver,
-                    jsonl_path.open(encoding="utf-8") as jsonl_file,
-                ):
-                    for line_number, line in enumerate(jsonl_file, start=1):
-                        if max_samples is not None and total >= max_samples:
-                            break
-                        if not line.strip():
-                            continue
-                        row = json.loads(line)
-                        sample_id = row.get("id")
-                        if not isinstance(sample_id, str) or not sample_id:
-                            raise ValueError(f"{jsonl_path}:{line_number}: row id must be a non-empty string.")
-                        key = _sample_key(subset, sample_id)
-                        if key in seen_keys:
-                            raise ValueError(f"{jsonl_path}:{line_number}: duplicate normalized sample key {key!r}.")
-                        seen_keys.add(key)
-                        messages, images = _normalize_messages(
-                            row.get("messages"),
-                            resolver=resolver,
-                            sample_name=f"{jsonl_path}:{line_number}",
+    writers: dict[str, wds.ShardWriter] = {}
+    with ExitStack() as writer_stack:
+        for subset in subsets:
+            subset_dir = (source_dir / subset).resolve()
+            if not subset_dir.is_relative_to(source_dir) or not subset_dir.is_dir():
+                raise FileNotFoundError(f"Downloaded subset directory does not exist: {subset_dir}")
+            jsonl_path = _jsonl_path(subset_dir, subset)
+            with (
+                MediaResolver(
+                    subset_dir,
+                    allow_loose_media=allow_loose_media,
+                    archive_paths=None if media_archives is None else media_archives.get(subset, ()),
+                ) as resolver,
+                jsonl_path.open(encoding="utf-8") as jsonl_file,
+            ):
+                for line_number, line in enumerate(jsonl_file, start=1):
+                    if max_samples is not None and total >= max_samples:
+                        break
+                    if not line.strip():
+                        continue
+                    row = json.loads(line)
+                    sample_id = row.get("id")
+                    if not isinstance(sample_id, str) or not sample_id:
+                        raise ValueError(f"{jsonl_path}:{line_number}: row id must be a non-empty string.")
+                    key = _sample_key(subset, sample_id)
+                    if key in seen_keys:
+                        raise ValueError(f"{jsonl_path}:{line_number}: duplicate normalized sample key {key!r}.")
+                    seen_keys.add(key)
+                    messages, images = _normalize_messages(
+                        row.get("messages"),
+                        resolver=resolver,
+                        sample_name=f"{jsonl_path}:{line_number}",
+                    )
+                    split = _sample_split(key, validation_fraction)
+                    if split not in writers:
+                        writers[split] = writer_stack.enter_context(
+                            _make_shard_writer(output_dir, split, max_samples_per_tar)
                         )
-                        split = _sample_split(key, validation_fraction)
-                        writers[split].write(key, images, messages)
-                        counts[split] += 1
-                        total += 1
-                if max_samples is not None and total >= max_samples:
-                    break
+                    _write_sample(writers[split], key, images, messages)
+                    counts[split] += 1
+                    total += 1
+            if max_samples is not None and total >= max_samples:
+                break
 
     if counts["train"] == 0:
         raise RuntimeError("No training samples were written.")