Add Megatron data manifest and g0050 setup
This commit is contained in:
340
dataset/pretrain/scripts/convert_pretrain_parquet_to_megatron.py
Executable file
340
dataset/pretrain/scripts/convert_pretrain_parquet_to_megatron.py
Executable file
@@ -0,0 +1,340 @@
|
||||
#!/usr/bin/env python3
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import sys
|
||||
import time
|
||||
from concurrent.futures import ProcessPoolExecutor, as_completed
|
||||
from multiprocessing import Pool
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
from typing import Iterable
|
||||
|
||||
import pyarrow.parquet as pq
|
||||
|
||||
_TOKENIZER = None
|
||||
_APPEND_EOD = True
|
||||
_MAX_SEQ_LEN = 65536
|
||||
|
||||
|
||||
def add_megatron_to_path(megatron_dir: str) -> None:
|
||||
path = str(Path(megatron_dir).resolve())
|
||||
if path not in sys.path:
|
||||
sys.path.insert(0, path)
|
||||
|
||||
|
||||
def make_tokenizer_args(args: argparse.Namespace) -> SimpleNamespace:
|
||||
return SimpleNamespace(
|
||||
rank=0,
|
||||
make_vocab_size_divisible_by=128,
|
||||
tensor_model_parallel_size=1,
|
||||
pad_vocab_size=True,
|
||||
padded_vocab_size=None,
|
||||
vocab_size=args.vocab_size,
|
||||
vocab_file=args.vocab_file,
|
||||
merge_file=args.merge_file,
|
||||
vocab_extra_ids=0,
|
||||
tokenizer_type=args.tokenizer_type,
|
||||
tokenizer_model=args.tokenizer_model,
|
||||
metadata_path=args.tokenizer_metadata,
|
||||
special_tokens=args.tokenizer_special_tokens,
|
||||
tokenizer_sentencepiece_legacy=args.tokenizer_sentencepiece_legacy,
|
||||
tokenizer_hf_no_use_fast=args.tokenizer_hf_no_use_fast,
|
||||
tokenizer_hf_no_include_special_tokens=args.tokenizer_hf_no_include_special_tokens,
|
||||
trust_remote_code=args.trust_remote_code,
|
||||
tiktoken_pattern=args.tiktoken_pattern,
|
||||
tiktoken_num_special_tokens=args.tiktoken_num_special_tokens,
|
||||
null_tokenizer_eod_id=args.null_tokenizer_eod_id,
|
||||
null_tokenizer_pad_id=args.null_tokenizer_pad_id,
|
||||
tokenizer_prompt_format=None,
|
||||
chat_template=None,
|
||||
image_tag_type=None,
|
||||
force_system_message=False,
|
||||
sft_tokenizer_prompt_format=None,
|
||||
)
|
||||
|
||||
|
||||
def build_tokenizer(args: argparse.Namespace):
|
||||
add_megatron_to_path(args.megatron_dir)
|
||||
from megatron.core.tokenizers.utils.build_tokenizer import build_tokenizer
|
||||
|
||||
return build_tokenizer(make_tokenizer_args(args))
|
||||
|
||||
|
||||
def init_worker(args: argparse.Namespace) -> None:
|
||||
global _TOKENIZER, _APPEND_EOD, _MAX_SEQ_LEN
|
||||
_APPEND_EOD = args.append_eod
|
||||
_MAX_SEQ_LEN = args.max_seq_len
|
||||
_TOKENIZER = build_tokenizer(args)
|
||||
if _APPEND_EOD and _TOKENIZER.eod is None:
|
||||
raise ValueError("Tokenizer has no EOD/EOS token, but --append-eod is enabled.")
|
||||
|
||||
|
||||
def encode_text(text: object):
|
||||
if not isinstance(text, str):
|
||||
return None
|
||||
text = text.strip()
|
||||
if not text:
|
||||
return None
|
||||
token_ids = _TOKENIZER.tokenize(text)
|
||||
if not token_ids:
|
||||
return None
|
||||
|
||||
if _MAX_SEQ_LEN and _MAX_SEQ_LEN > 0:
|
||||
content_chunk_size = _MAX_SEQ_LEN - 1 if _APPEND_EOD else _MAX_SEQ_LEN
|
||||
chunks = []
|
||||
for start in range(0, len(token_ids), content_chunk_size):
|
||||
chunk = token_ids[start : start + content_chunk_size]
|
||||
if _APPEND_EOD:
|
||||
chunk = [*chunk, _TOKENIZER.eod]
|
||||
chunks.append((chunk, [len(chunk)]))
|
||||
return chunks
|
||||
|
||||
if _APPEND_EOD:
|
||||
token_ids.append(_TOKENIZER.eod)
|
||||
return [(token_ids, [len(token_ids)])]
|
||||
|
||||
|
||||
def output_paths(output_prefix: Path, text_key: str) -> tuple[Path, Path]:
|
||||
return (
|
||||
Path(str(output_prefix) + f"_{text_key}_document.bin"),
|
||||
Path(str(output_prefix) + f"_{text_key}_document.idx"),
|
||||
)
|
||||
|
||||
|
||||
def resolve_inputs(inputs: Iterable[str], max_files: int | None) -> list[Path]:
|
||||
paths: list[Path] = []
|
||||
for item in inputs:
|
||||
path = Path(item)
|
||||
if path.is_dir():
|
||||
paths.extend(path.glob("*.parquet"))
|
||||
paths.extend(path.glob("*.zstd.parquet"))
|
||||
continue
|
||||
matches = sorted(Path().glob(item)) if not path.is_absolute() else sorted(path.parent.glob(path.name))
|
||||
if matches:
|
||||
paths.extend(matches)
|
||||
continue
|
||||
if path.exists():
|
||||
paths.append(path)
|
||||
continue
|
||||
raise FileNotFoundError(item)
|
||||
unique = sorted({p.resolve() for p in paths}, key=str)
|
||||
if max_files is not None:
|
||||
unique = unique[:max_files]
|
||||
return unique
|
||||
|
||||
|
||||
def safe_prefix_component(value: str) -> str:
|
||||
value = re.sub(r"[^A-Za-z0-9_.-]+", "_", value.strip())
|
||||
value = value.strip("._-")
|
||||
return value or "source"
|
||||
|
||||
|
||||
def convert_one_parquet(task: tuple[str, argparse.Namespace]) -> dict[str, object]:
|
||||
parquet_path = Path(task[0])
|
||||
args = task[1]
|
||||
add_megatron_to_path(args.megatron_dir)
|
||||
from megatron.core.datasets import indexed_dataset
|
||||
|
||||
stem = parquet_path.name.replace(".zstd.parquet", "").replace(".parquet", "")
|
||||
source_name = safe_prefix_component(parquet_path.parent.name)
|
||||
output_prefix = Path(args.output_dir) / f"{args.output_prefix_prefix}_{source_name}_{stem}"
|
||||
bin_file, idx_file = output_paths(output_prefix, args.text_key)
|
||||
if not args.overwrite and bin_file.exists() and idx_file.exists():
|
||||
return {
|
||||
"status": "skip",
|
||||
"source": str(parquet_path),
|
||||
"output_prefix": str(output_prefix),
|
||||
"reason": "existing bin/idx",
|
||||
}
|
||||
|
||||
bin_file.unlink(missing_ok=True)
|
||||
idx_file.unlink(missing_ok=True)
|
||||
output_prefix.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
tokenizer = build_tokenizer(args)
|
||||
dtype = indexed_dataset.DType.optimal_dtype(tokenizer.vocab_size)
|
||||
builder = indexed_dataset.IndexedDatasetBuilder(str(bin_file), dtype=dtype)
|
||||
|
||||
start = time.time()
|
||||
rows = source_docs = output_docs = tokens = split_source_docs = split_extra_docs = 0
|
||||
stop = False
|
||||
|
||||
def consume(encoded) -> None:
|
||||
nonlocal source_docs, output_docs, tokens, split_source_docs, split_extra_docs, stop
|
||||
if encoded is None or stop:
|
||||
return
|
||||
source_docs += 1
|
||||
if len(encoded) > 1:
|
||||
split_source_docs += 1
|
||||
split_extra_docs += len(encoded) - 1
|
||||
for token_ids, lengths in encoded:
|
||||
if args.max_seq_len and len(token_ids) > args.max_seq_len:
|
||||
raise ValueError(
|
||||
f"internal split error: chunk has {len(token_ids)} tokens, "
|
||||
f"max_seq_len={args.max_seq_len}"
|
||||
)
|
||||
builder.add_document(token_ids, lengths)
|
||||
output_docs += 1
|
||||
tokens += len(token_ids)
|
||||
if args.log_interval and source_docs % args.log_interval == 0:
|
||||
elapsed = max(time.time() - start, 1e-6)
|
||||
print(
|
||||
f"[{parquet_path.name}] source_docs={source_docs} output_docs={output_docs} "
|
||||
f"tokens={tokens} source_docs/s={source_docs / elapsed:.2f}",
|
||||
flush=True,
|
||||
)
|
||||
if args.max_docs and source_docs >= args.max_docs:
|
||||
stop = True
|
||||
|
||||
pf = pq.ParquetFile(parquet_path)
|
||||
if args.text_key not in pf.schema_arrow.names:
|
||||
raise ValueError(f"{parquet_path} has no column {args.text_key!r}")
|
||||
|
||||
if args.workers_per_file == 1:
|
||||
init_worker(args)
|
||||
for batch in pf.iter_batches(columns=[args.text_key], batch_size=args.batch_size, use_threads=True):
|
||||
texts = batch.column(0).to_pylist()
|
||||
rows += len(texts)
|
||||
for text in texts:
|
||||
consume(encode_text(text))
|
||||
if stop:
|
||||
break
|
||||
if stop:
|
||||
break
|
||||
else:
|
||||
with Pool(processes=args.workers_per_file, initializer=init_worker, initargs=(args,)) as pool:
|
||||
for batch in pf.iter_batches(columns=[args.text_key], batch_size=args.batch_size, use_threads=True):
|
||||
texts = batch.column(0).to_pylist()
|
||||
rows += len(texts)
|
||||
for encoded in pool.imap(encode_text, texts, chunksize=args.chunksize):
|
||||
consume(encoded)
|
||||
if stop:
|
||||
break
|
||||
if stop:
|
||||
break
|
||||
|
||||
builder.finalize(str(idx_file))
|
||||
elapsed = max(time.time() - start, 1e-6)
|
||||
return {
|
||||
"status": "ok",
|
||||
"source": str(parquet_path),
|
||||
"output_prefix": str(output_prefix),
|
||||
"bin_file": str(bin_file),
|
||||
"idx_file": str(idx_file),
|
||||
"dtype": dtype.__name__,
|
||||
"rows": rows,
|
||||
"source_docs": source_docs,
|
||||
"docs": output_docs,
|
||||
"output_docs": output_docs,
|
||||
"split_source_docs": split_source_docs,
|
||||
"split_extra_docs": split_extra_docs,
|
||||
"max_seq_len": args.max_seq_len,
|
||||
"tokens": tokens,
|
||||
"elapsed_sec": elapsed,
|
||||
"source_docs_per_sec": source_docs / elapsed,
|
||||
"docs_per_sec": output_docs / elapsed,
|
||||
}
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description="Convert pretrain parquet directly to Megatron .bin/.idx.")
|
||||
parser.add_argument("--input", action="append", required=True, help="Parquet dir/file/glob. Repeatable.")
|
||||
parser.add_argument("--output-dir", required=True)
|
||||
parser.add_argument("--manifest", default=None)
|
||||
parser.add_argument("--megatron-dir", default="/opt/Megatron-Bridge/3rdparty/Megatron-LM")
|
||||
parser.add_argument("--tokenizer-type", default="HuggingFaceTokenizer")
|
||||
parser.add_argument("--tokenizer-model", required=True)
|
||||
parser.add_argument("--tokenizer-metadata", default=None)
|
||||
parser.add_argument("--tokenizer-special-tokens", nargs="*", default=None)
|
||||
parser.add_argument("--tokenizer-sentencepiece-legacy", action="store_true")
|
||||
parser.add_argument("--tokenizer-hf-no-use-fast", action="store_true")
|
||||
parser.add_argument("--tokenizer-hf-no-include-special-tokens", action="store_true")
|
||||
parser.add_argument("--trust-remote-code", action="store_true")
|
||||
parser.add_argument("--vocab-file", default=None)
|
||||
parser.add_argument("--merge-file", default=None)
|
||||
parser.add_argument("--vocab-size", type=int, default=None)
|
||||
parser.add_argument("--tiktoken-pattern", default=None)
|
||||
parser.add_argument("--tiktoken-num-special-tokens", type=int, default=1000)
|
||||
parser.add_argument("--null-tokenizer-eod-id", type=int, default=None)
|
||||
parser.add_argument("--null-tokenizer-pad-id", type=int, default=-1)
|
||||
parser.add_argument("--text-key", default="text")
|
||||
parser.add_argument("--output-prefix-prefix", default="laoyao_2b_moe_8192")
|
||||
parser.add_argument("--parallel-files", type=int, default=1)
|
||||
parser.add_argument("--workers-per-file", type=int, default=max((os.cpu_count() or 8) // 4, 1))
|
||||
parser.add_argument("--batch-size", type=int, default=8192)
|
||||
parser.add_argument("--chunksize", type=int, default=128)
|
||||
parser.add_argument("--max-files", type=int, default=None)
|
||||
parser.add_argument("--max-docs", type=int, default=0, help="Per-file doc cap; 0 means no cap.")
|
||||
parser.add_argument("--max-seq-len", type=int, default=65536, help="Split documents into chunks no longer than this; 0 disables splitting.")
|
||||
parser.add_argument("--log-interval", type=int, default=10000)
|
||||
parser.add_argument("--append-eod", action=argparse.BooleanOptionalAction, default=True)
|
||||
parser.add_argument("--overwrite", action="store_true")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
if args.parallel_files < 1 or args.workers_per_file < 1:
|
||||
raise ValueError("parallel and worker counts must be positive")
|
||||
if args.max_seq_len < 0:
|
||||
raise ValueError("--max-seq-len must be >= 0")
|
||||
if args.append_eod and args.max_seq_len == 1:
|
||||
raise ValueError("--max-seq-len must be > 1 when --append-eod is enabled")
|
||||
files = resolve_inputs(args.input, args.max_files)
|
||||
if not files:
|
||||
raise FileNotFoundError("no parquet files resolved")
|
||||
|
||||
Path(args.output_dir).mkdir(parents=True, exist_ok=True)
|
||||
manifest_path = Path(args.manifest) if args.manifest else Path(args.output_dir) / "manifest.json"
|
||||
print(
|
||||
json.dumps(
|
||||
{
|
||||
"files": len(files),
|
||||
"parallel_files": args.parallel_files,
|
||||
"workers_per_file": args.workers_per_file,
|
||||
"total_tokenizer_workers": args.parallel_files * args.workers_per_file,
|
||||
"output_dir": args.output_dir,
|
||||
"max_seq_len": args.max_seq_len,
|
||||
},
|
||||
ensure_ascii=False,
|
||||
),
|
||||
flush=True,
|
||||
)
|
||||
|
||||
tasks = [(str(path), args) for path in files]
|
||||
results: list[dict[str, object]] = []
|
||||
if args.parallel_files == 1:
|
||||
for task in tasks:
|
||||
result = convert_one_parquet(task)
|
||||
results.append(result)
|
||||
print(json.dumps(result, ensure_ascii=False), flush=True)
|
||||
else:
|
||||
with ProcessPoolExecutor(max_workers=args.parallel_files) as executor:
|
||||
futures = [executor.submit(convert_one_parquet, task) for task in tasks]
|
||||
for future in as_completed(futures):
|
||||
result = future.result()
|
||||
results.append(result)
|
||||
print(json.dumps(result, ensure_ascii=False), flush=True)
|
||||
|
||||
manifest = {
|
||||
"args": vars(args),
|
||||
"files": [str(path) for path in files],
|
||||
"results": sorted(results, key=lambda item: str(item.get("source", ""))),
|
||||
"total_source_docs": sum(int(r.get("source_docs", 0)) for r in results),
|
||||
"total_docs": sum(int(r.get("docs", 0)) for r in results),
|
||||
"total_output_docs": sum(int(r.get("output_docs", 0)) for r in results),
|
||||
"total_split_source_docs": sum(int(r.get("split_source_docs", 0)) for r in results),
|
||||
"total_split_extra_docs": sum(int(r.get("split_extra_docs", 0)) for r in results),
|
||||
"total_tokens": sum(int(r.get("tokens", 0)) for r in results),
|
||||
"ok_prefixes": [str(r["output_prefix"]) for r in results if r.get("status") in {"ok", "skip"}],
|
||||
}
|
||||
manifest_path.write_text(json.dumps(manifest, ensure_ascii=False, indent=2), encoding="utf-8")
|
||||
print(f"manifest={manifest_path}", flush=True)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user