[Feature] Add Nemotron-CC v1 and v2 dataset support - #83
Conversation
akaashrp
changed the title
akaashrp
changed the title
|
CI never ran here. GitHub's Actions incident (Aug 6, 15:05–00:14 UTC) throttled webhooks to ~15%, so the |
|
@claude review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| scheduler.load_state_dict(st) | ||
|
|
||
|
|
||
| def maybe_reload_dataset_mixture(cfg: PretrainLMCfg) -> None: |
There was a problem hiding this comment.
Compactness: The entire online-reweighting path is a speculative hook — no training entrypoint in the repo sets dataset_mixture_hot_reload_path. This adds ~65 lines here (mtime polling, per-step broadcast, versioning) plus two config knobs (dataset_mixture_hot_reload_path, dataset_mixture_poll_interval_steps) and two context fields (dataset_mixture_last_mtime_ns, dataset_mixture_version), all for a feature nothing exercises.
Simpler: set the mixture once in setup_dataset from cfg.dataset_mixture, and drop this function, its train_step call, the two config fields, and the two context fields. Add the reload path back when a run actually drives it.
| self._weights: Dict[str, float] = {} | ||
| self._weight_vec = np.empty(0, dtype=np.float64) |
There was a problem hiding this comment.
Compactness: _weight_vec is written here and in update_weights (line 206) but never read anywhere. Dead field — drop the initializer here and the assignment in update_weights.
| self._weights: Dict[str, float] = {} | |
| self._weight_vec = np.empty(0, dtype=np.float64) | |
| self._weights: Dict[str, float] = {} | |
| self._cdf = np.empty(0, dtype=np.float64) |
| if __name__ == "__main__": | ||
| raw_root = Path("workspace/datasets/nemotron-cc-v1/rawtxt") | ||
| download_raw_files(raw_root) | ||
|
|
||
| if __name__ == "__main__": | ||
| cfg = TokenizeCorpusCfg() | ||
| cfg.tokenizer_name = "Qwen/Qwen3-30B-A3B" | ||
| cfg.source_path = Path("workspace/datasets/nemotron-cc-v1/rawtxt") | ||
| cfg.output_path = Path("workspace/datasets/nemotron-cc-v1/toktxt/qwen3") | ||
| launch(cfg) |
There was a problem hiding this comment.
Compactness: Two consecutive if __name__ == "__main__": guards where one does. Merge them (and reuse raw_root):
| if __name__ == "__main__": | |
| raw_root = Path("workspace/datasets/nemotron-cc-v1/rawtxt") | |
| download_raw_files(raw_root) | |
| if __name__ == "__main__": | |
| cfg = TokenizeCorpusCfg() | |
| cfg.tokenizer_name = "Qwen/Qwen3-30B-A3B" | |
| cfg.source_path = Path("workspace/datasets/nemotron-cc-v1/rawtxt") | |
| cfg.output_path = Path("workspace/datasets/nemotron-cc-v1/toktxt/qwen3") | |
| launch(cfg) | |
| if __name__ == "__main__": | |
| raw_root = Path("workspace/datasets/nemotron-cc-v1/rawtxt") | |
| download_raw_files(raw_root) | |
| cfg = TokenizeCorpusCfg() | |
| cfg.tokenizer_name = "Qwen/Qwen3-30B-A3B" | |
| cfg.source_path = raw_root | |
| cfg.output_path = Path("workspace/datasets/nemotron-cc-v1/toktxt/qwen3") | |
| launch(cfg) |
| if __name__ == "__main__": | ||
| raw_root = Path("workspace/datasets/nemotron-cc-v2/rawtxt") | ||
| download_raw_files(raw_root) | ||
|
|
||
| if __name__ == "__main__": | ||
| cfg = TokenizeCorpusCfg() | ||
| cfg.tokenizer_name = "Qwen/Qwen3-30B-A3B" | ||
| cfg.source_path = Path("workspace/datasets/nemotron-cc-v2/rawtxt") | ||
| cfg.output_path = Path("workspace/datasets/nemotron-cc-v2/toktxt/qwen3") | ||
| launch(cfg) |
There was a problem hiding this comment.
Compactness: Same doubled if __name__ == "__main__": guard as v1. Merge into one block:
| if __name__ == "__main__": | |
| raw_root = Path("workspace/datasets/nemotron-cc-v2/rawtxt") | |
| download_raw_files(raw_root) | |
| if __name__ == "__main__": | |
| cfg = TokenizeCorpusCfg() | |
| cfg.tokenizer_name = "Qwen/Qwen3-30B-A3B" | |
| cfg.source_path = Path("workspace/datasets/nemotron-cc-v2/rawtxt") | |
| cfg.output_path = Path("workspace/datasets/nemotron-cc-v2/toktxt/qwen3") | |
| launch(cfg) | |
| if __name__ == "__main__": | |
| raw_root = Path("workspace/datasets/nemotron-cc-v2/rawtxt") | |
| download_raw_files(raw_root) | |
| cfg = TokenizeCorpusCfg() | |
| cfg.tokenizer_name = "Qwen/Qwen3-30B-A3B" | |
| cfg.source_path = raw_root | |
| cfg.output_path = Path("workspace/datasets/nemotron-cc-v2/toktxt/qwen3") | |
| launch(cfg) |
Compactness ReviewThe core of this PR — parquet reading,
Details in the inline threads. |
| object_list = [payload] | ||
| torch.distributed.broadcast_object_list(object_list, src=0) |
There was a problem hiding this comment.
Performance: broadcast_object_list runs on every poll, unconditionally, even when the mixture is unchanged (payload=None). With the default dataset_mixture_poll_interval_steps=1, that's a pickle-based world collective on the hot train_step path every step whenever hot reload is enabled. On NCCL, broadcast_object_list stages the pickled bytes through a GPU tensor and copies back to CPU to unpickle — a device sync per step that serializes against pipeline compute.
Gate the expensive broadcast behind a cheap check: rank 0 broadcasts a single 1-element int tensor (the mixture version, or a changed flag); only when it advances do all ranks do the object broadcast. That keeps the steady-state per-step cost to one small tensor broadcast instead of a pickle round-trip + sync.
| dataset_mixture_hot_reload_path: Optional[Path] = None | ||
| """Optional JSON file path for online mixture reweighting.""" | ||
|
|
||
| dataset_mixture_poll_interval_steps: int = 1 |
There was a problem hiding this comment.
Performance: Default of 1 means the hot-reload machinery in maybe_reload_dataset_mixture fires every single step (a stat() on rank 0 + a world broadcast_object_list — see the thread there). Mixture reweighting is an operator-driven, coarse-grained action; a per-step poll buys nothing and adds a per-step collective. Default to something coarse (e.g. hundreds of steps) so the steady-state training loop isn't paying for it every step.
Consistency ReviewDocs and comments track the code changes cleanly. The No stale reference was found: grepping AGENTS.md, docs/, and .agents/skills/ for the renamed/changed surfaces (input formats, |
Performance ReviewThis PR is data-loading/tokenization plumbing — it doesn't touch the GPU hot path (DualPipeV overlap, kernels, compiled regions), so no risk to comm overlap or graph stability. Two things to address, both scoped to the new weighted-mixture path (the default single-source
The new per-sample splitmix64 sampling in Evidence gap: the PR body is empty and no throughput/step-time numbers are provided. Since it adds a per-step call into the training loop, a tokens-per-second comparison of single-source vs. mixture mode (and mixture with hot reload enabled, at the default poll interval) would confirm the steady-state loop doesn't regress. |
| training.dataset_mixture_last_mtime_ns = None | ||
| training.dataset_mixture_version = 0 | ||
| if cfg.dataset_sources: | ||
| if not cfg.dataset_mixture: | ||
| raise ValueError("dataset_mixture must be provided when dataset_sources is set.") | ||
| source_datasets: dict[str, SourceDataset] = {} | ||
| for source_name, source_root in sorted(cfg.dataset_sources.items()): | ||
| memmap_datasets = [] | ||
| for file in sorted(source_root.rglob("*.bin")): | ||
| memmap_datasets.append(MemmapDataset(file, cfg.sequence_length)) | ||
| source_datasets[source_name] = SourceDataset(source_name, memmap_datasets) | ||
| training.dataset = WeightedMixtureDataset(source_datasets, cfg.seed, cfg.dataset_mixture) |
There was a problem hiding this comment.
Correctness: Hot-reloaded mixture state is not part of the checkpoint, so resume silently trains on the wrong data distribution. AppState/load_checkpoint persist only model/optim/scheduler/rng/step; on resume setup_dataset resets dataset_mixture_version = 0, dataset_mixture_last_mtime_ns = None, and rebuilds the mixture from the initial cfg.dataset_mixture, discarding every hot-reload applied before the checkpoint.
Concrete failure: a curriculum applies weights W2 via the hot-reload file at step 1000, checkpoint at step 2000 with W2 active, crash, resume. If the hot-reload file was since removed/renamed (or dataset_mixture_poll_interval_steps > 1 and the resumed step isn't a poll multiple), mtime_ns is None/non-triggering, so W2 is never re-applied and training continues on the initial W0 mixture. Even with the default poll_interval == 1, a deleted file makes the revert permanent.
To settle it, persist dataset_mixture_version and the current weights in the checkpoint (or a sidecar) and restore them in load_checkpoint before the first maybe_reload_dataset_mixture, so the run continues from the last-applied mixture regardless of the file's current state.
Correctness ReviewThis PR is data-pipeline-only — Nemotron-CC v1/v2 fetch + parquet reading, and a new I traced the sampling and the hot-reload collective and both look correct: One concrete correctness gap: hot-reloaded mixture weights and version are not checkpointed, so a resume can silently revert to the initial config mixture (details + a fix direction in the inline thread on |
e5bd465 to
ca327ca
Compare
ca327ca to
7ca6cdc
Compare
2 participants
No description provided.