diff --git a/small_llm_pretraining/nemo/Dockerfile.h200 b/small_llm_pretraining/nemo/Dockerfile.h200 new file mode 100644 index 000000000..5cf32b704 --- /dev/null +++ b/small_llm_pretraining/nemo/Dockerfile.h200 @@ -0,0 +1,134 @@ +# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a +# copy of this software and associated documentation files (the "Software"), +# to deal in the Software without restriction, including without limitation +# the rights to use, copy, modify, merge, publish, distribute, sublicense, +# and/or sell copies of the Software, and to permit persons to whom the +# Software is furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL +# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING +# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER +# DEALINGS IN THE SOFTWARE. + + +ARG FROM_IMAGE_NAME=nvcr.io/nvidia/pytorch:25.01-py3 +FROM ${FROM_IMAGE_NAME} + +# Document build setup +ARG FROM_IMAGE_NAME +ENV CUSTOM_FROM_IMAGE_NAME ${FROM_IMAGE_NAME} + +# Custom libraries version +WORKDIR /workspace/ + +ARG GIT_COMMIT_ID +ENV GIT_COMMIT_ID=$GIT_COMMIT_ID + +RUN git config --global user.name "a" && \ + git config --global user.email "a" + +WORKDIR /workspace/ + +RUN pip install numcodecs==0.13.1 + +## 1. Apex +ARG APEX_REVISION=SKIP +ENV CUSTOM_APEX_REVISION ${APEX_REVISION} +ARG APEX_MAX_JOBS=4 + +RUN if [ "${APEX_REVISION}" != SKIP ]; then \ + git clone https://github.com/NVIDIA/apex && \ + cd apex && \ + echo APEX_REVISION=${APEX_REVISION} && \ + git checkout ${APEX_REVISION} && \ + echo APEX_COMMIT_HASH=$(git rev-parse HEAD) && \ + MAX_JOBS=${APEX_MAX_JOBS} NVCC_APPEND_FLAGS="--threads 8" pip install -v --no-build-isolation --no-cache-dir --disable-pip-version-check --config-settings "--build-option=--cpp_ext --cuda_ext --bnp --xentropy --deprecated_fused_adam --deprecated_fused_lamb --fast_multihead_attn --distributed_lamb --fast_layer_norm --transducer --distributed_adam --fmha --fast_bottleneck --nccl_p2p --peer_memory --permutation_search --focal_loss --fused_conv_bias_relu --index_mul_2d --cudnn_gbn --group_norm" . \ + ; fi + + + +## 2. Transformer Engine +ARG TE_REVISION=SKIP +ENV CUSTOM_TE_REVISION ${TE_REVISION} + +RUN if [ "${TE_REVISION}" != SKIP ]; then \ + pip uninstall -y transformer-engine && \ + git clone https://github.com/NVIDIA/TransformerEngine.git transformerengine && \ + cd transformerengine && \ + git checkout ${TE_REVISION} && \ + echo TE_COMMIT_HASH=$(git rev-parse HEAD) && \ + echo $(git rev-parse HEAD) > /TE_COMMIT_HASH.env && \ + git submodule init && git submodule update && \ + NVTE_CUDA_ARCHS="90;100" NVTE_UB_WITH_MPI=1 NVTE_FRAMEWORK=pytorch MPI_HOME=/usr/local/mpi pip install --force-reinstall --no-deps . \ + ; fi + + +## 3. NeMo +ARG NEMO_REVISION=v2.1.0 +ENV CUSTOM_NEMO_REVISION ${NEMO_REVISION} + +# Clone and checkout NeMo at specified version +RUN git clone https://github.com/NVIDIA/NeMo.git && \ + cd NeMo && \ + git checkout ${NEMO_REVISION} && \ + echo NEMO_COMMIT_HASH=$(git rev-parse HEAD) && \ + echo $(git rev-parse HEAD) > /NEMO_COMMIT_HASH.env && \ + pip uninstall -y nemo-toolkit sacrebleu && \ + # Only keep edits that are necessary (remove AMD-specific workarounds if upstream doesn't need them) + sed -i "/mamba-ssm/d" requirements/requirements_nlp.txt && \ + sed -i 's/tensorstore<0.1.46/tensorstore/g' requirements/requirements_nlp.txt && \ + sed -i 's/protobuf==3.20.3/protobuf/g' requirements/requirements.txt && \ + pip install "cython<3.0.0" && \ + pip install -e ".[llm]" && \ + pip install -e ".[nlp]" + + +## 3.1 NeMo-Run +ARG NEMORUN_REVISION=v0.4.0 +ENV CUSTOM_NEMORUN_REVISION ${NEMORUN_REVISION} + +RUN git clone https://github.com/NVIDIA/NeMo-Run.git && \ + cd NeMo-Run && \ + git checkout ${NEMORUN_REVISION} && \ + echo NEMORUN_COMMIT_HASH=$(git rev-parse HEAD) && \ + pip install -e . + + +# Python deps +# Important this should be done after NeMo, otherwise the pinned transformers==4.40.2 version will be overwritten +COPY requirements.txt requirements.txt +RUN pip3 install -r requirements.txt + + +# 4. Megatron-core +ARG MCORE_REVISION=core_r0.11.0 +ARG MCORE_REPO=https://github.com/NVIDIA/Megatron-LM.git +ENV CUSTOM_MCORE_REVISION ${MCORE_REVISION} + +RUN if [ "${MCORE_REVISION}" != SKIP ]; then \ + pip uninstall -y megatron-core && \ + git clone ${MCORE_REPO} Megatron-LM && \ + cd Megatron-LM && \ + git checkout ${MCORE_REVISION} && \ + echo MCORE_COMMIT_HASH=$(git rev-parse HEAD) && \ + echo $(git rev-parse HEAD) > /MCORE_COMMIT_HASH.env && \ + pip install . && \ + cd megatron/core/datasets && \ + make \ + ; fi + +ENV PYTHONPATH "${PYTHONPATH}:/workspace/Megatron-LM" + + +WORKDIR /workspace/code + +# Copy the current state of the code inside the image +COPY . . \ No newline at end of file diff --git a/small_llm_pretraining/nemo/Dockerfile.mi325 b/small_llm_pretraining/nemo/Dockerfile.mi325 new file mode 100644 index 000000000..1d0f7be82 --- /dev/null +++ b/small_llm_pretraining/nemo/Dockerfile.mi325 @@ -0,0 +1,93 @@ +FROM rocm/pytorch:rocm6.4_ubuntu22.04_py3.10_pytorch_release_2.6.0 + +WORKDIR /workspace + +RUN pip install pybind11 +RUN pip install ninja +RUN pip install packaging +RUN /usr/bin/python3 -m pip install pyYAML + +# Install library dependencies +WORKDIR /workspace/deps + +# FlashAttention +RUN git clone https://github.com/ROCm/flash-attention/ flash_attention \ + # latest stable commit of ck_tile/fa3 branch + && cd flash_attention && git checkout cace3592812640486b04196a209bb85d12267b4c \ + && git submodule update --init --recursive \ + && PYTORCH_ROCM_ARCH='gfx942' GPU_ARCHS="gfx942" MAX_JOBS=64 pip install --no-build-isolation -e . + +ADD patches /workspace/deps/patches + +# Megatron-core +RUN git clone --recursive https://github.com/ROCm/Megatron-LM.git megatron_lm +RUN pip uninstall -y megatron-core +# dev branch commit +RUN cd megatron_lm && git checkout megatron_190213a_mlperf \ + && pip install -e . && cd megatron/core/datasets && make + +ENV PYTHONPATH "${PYTHONPATH}:/workspace/deps/megatron_lm" + +# mambe dependency required for NeMo +RUN git clone https://github.com/state-spaces/mamba.git mamba_ssm \ + && cd mamba_ssm \ + && git checkout v2.2.2 \ + && export HIP_ARCHITECTURES="gfx942" \ + && pip install --no-cache-dir --verbose . + +# NeMo +RUN git clone https://github.com/NVIDIA/NeMo nemo \ + && cd nemo && git checkout v2.1.0 +RUN cd /workspace/deps/nemo \ + && git apply /workspace/deps/patches/nemo_v2_1_0.patch \ + && pip install --no-build-isolation -e ".[nlp]" + +# NeMo-Run +RUN pip install git+https://github.com/NVIDIA/NeMo-Run.git@v0.4.0 + +# Python deps +# Important this should be done after NeMo, otherwise the pinned transformers==4.40.2 version will be overwritten +COPY requirements.txt requirements.txt +RUN pip3 install -r requirements.txt + +# Transformer Engine +ARG TE_COMMIT=te_v1.9_mlperf_llama2 +RUN git clone --recursive https://github.com/ROCm/TransformerEngine.git \ + # dev branch commit + && cd TransformerEngine && git checkout $TE_COMMIT && git submodule update --init --recursive \ + # Workaround logging debug info to the console + && sed -i 's/self.logger.info/self.logger.debug/g' /workspace/deps/TransformerEngine/transformer_engine/pytorch/attention.py \ + && sed -i 's/warnings.warn/if False: warnings.warn/g' /workspace/deps/TransformerEngine/transformer_engine/pytorch/attention.py \ + && sed -i '/.*\"window_size should be.*/d' /workspace/deps/TransformerEngine/transformer_engine/common/fused_attn_rocm/fused_attn.cpp \ + && NVTE_FUSED_ATTN_AOTRITON=0 NVTE_ROCM_ARCH='gfx942' NVTE_FRAMEWORK='pytorch' NVTE_USE_HIPBLASLT=1 MAX_JOBS=128 PYTORCH_ROCM_ARCH='gfx942' GPU_ARCHS='gfx942' pip install -e . + +# Install hipBLASLt (FP8 tuned gemms - second round) + RUN git clone https://github.com/ROCm/hipBLASLt.git \ + && cd hipBLASLt && git checkout ebc770851dfb99a1bbb8ef2e5873c753f8011a47 \ + && sudo apt-get update \ + && apt install -y python3.10-venv \ + && ./install.sh -idc -a gfx942 + +# RPD +RUN sudo apt-get update && \ + apt --fix-broken install -y && \ + apt-get install -y\ + sqlite3 libsqlite3-dev \ + libfmt-dev + +RUN git clone https://github.com/ROCmSoftwarePlatform/rocmProfileData \ + && cd rocmProfileData \ + && cd rocpd_python \ + && python3 setup.py bdist_wheel \ + && pip install dist/*.whl \ + && cd .. \ + && cd rpd_tracer \ + && python3 setup.py bdist_wheel \ + && pip3 install dist/*.whl \ + && cd .. \ + && make; make install + +WORKDIR /workspace/code + +# Copy the current state of the code inside the image +COPY . . \ No newline at end of file diff --git a/small_llm_pretraining/nemo/README.md b/small_llm_pretraining/nemo/README.md new file mode 100644 index 000000000..2c967751a --- /dev/null +++ b/small_llm_pretraining/nemo/README.md @@ -0,0 +1,210 @@ +# 1. Problem + +Small Language Model pretraining - Llama 3.1 8B + +# 2. Directions + + +#### Container setup + +To build the container: + +```bash +docker build -t -f Dockerfile . +``` + +To launch the container: + +``` +bash dev/run_docker.sh +``` + +### Steps to download and verify data + +The current codebase is using C4 dataset for train and evaluation. Please refer to [Section 3](#preprocessed-data-download) for downloading the preprocessed dataset and [Section 6](#data-preprocessing) if you would like to perform manual tokenization. + +### Steps to run and time + +To train Llama 3.1 8B, we need to fill out all fields in [config.sh](./config.sh). This file contains all configurations for Slurm cluster access and job submission configurations, directory mappings, containers, and model configurations. + +Once the `config.sh` is properly filled, we run the following code snippet **inside the container**: + +```bash +source config.sh +bash run_llama31.sh +``` + +# 3. Dataset/Environment +### Publication/Attribution + +We use the c4/en/3.0.1 dataset from [HuggingFace/AllenAI](https://huggingface.co/datasets/allenai/c4). + +### Preprocessed data download + +The pre-tokenized dataset and the tokenizer are available to download from the S3 bucket. You can download this data from the bucket using RClone as follows: + +To run Rclone on Windows, you can download the executable here. To install Rclone on Linux/macOS/BSD systems, run: + +``` +sudo -v ; curl https://rclone.org/install.sh | sudo bash +``` + +Once Rclone is installed, run the following command to authenticate with the bucket: + +``` +rclone config create mlc-training s3 provider=Cloudflare access_key_id=76ea42eadb867e854061a1806220ee1e secret_access_key=a53625c4d45e3ca8ac0df8a353ea3a41ffc3292aa25259addd8b7dc5a6ce2936 endpoint=https://c2686074cb2caf5cbaf6d134bdba8b47.r2.cloudflarestorage.com +``` + +You can then navigate in the terminal to your desired download directory and run the following commands to download the dataset and checkpoints: + +### Raw data downloading + +We use [AllenAI C4](https://huggingface.co/datasets/allenai/c4) dataset for this benchmark. The original zipped **`json.gz`** files can be downloaded by following AllenAI C4's instruction, and you can download our zipped customized validation dataset from the MLCommons S3 bucket by running the following command: + + +```bash +export ORIGINAL_C4_PATH="" + +# download the full C4 files, including all raw train and validations +rclone copy mlc-training:mlcommons-training-wg-public/common/datasets/c4/original/en_json/3.0.1 $ORIGINAL_C4_PATH -P +``` +After the download is complete, you should see files with the following naming conventions under `PREPROCESSED_PATH`, ending with both `.idx` and `.bin`: +- Training partitions: `c4-train.en__text_document` +- Validation partitions: `c4-validation-91205-samples.en_text_document` + +### Run data preprocessing + +After downloading, run the following command to process them to zip them into `.gz` format before running the data preprocessing. + +``` +bash utils/parallel_compress_json_to_gz.sh +``` + + +Run the following commands to merge all 1024 training files into 8 `json.gz` files, all 8 validation files into a single `json.gz` file, as well as generate our customized validation dataset. Each of the `json.gz` files will subsequently be preprocessed into a pair of megatron dataset files (`.bin` and `.idx`) by our preprocess.sh script. + +```bash +export C4_PATH="" +export MERGED_C4_PATH="" +# more information about this knob can be found in consolidate_data.sh +export N_VALIDATION_SAMPLES=91205 + +bash utils/consolidate_data.sh +``` + + +### Tokenizer + +We are using the Llama 3.1 8B tokenizer. You can run `utils/download_hf_llama3.sh` to download it. + + +After the data consolidation is done, we can run this [script](./utils/preprocess.sh) to perform preprocessing. To run the preprocessing script, we need to use the following commands: + +```bash +# fill in the built container path here +export CONT_IMAGE_URL="" +# pass in the folder path that contains the Llama tokenizer here +# please refer to the tokenizer section above for more details +export TOKENIZER_PATH="" +# pass in the merged file path here +export MERGED_C4_PATH="" +# this path is used for storing the preprocessed .bin and .idx files +export PREPROCESSED_PATH="" + +# Extra Slurm-related arguments can be provided here +sbatch utils/preprocess.sh +``` + +If you are not using Slurm, then you should go inside the `utils/preprocess.sh` and run the commands manually. + +Warning! If you receive an error message of file not found, look into where `preprocess_data_for_megatron.py` is located in your path. + + +#### Training and test data separation + +We use the default split from the C4 dataset. This means that we use `c4-train.-of-01024.json.gz` files (where `768 <= x <= 1023`) for training, and we use our customized `c4-validation-91205-samples.en.json.gz`, which contains the first 91205 samples from the unshuffled C4 validation dataset, for evaluation. + +Notice here that we are using the first 1024 sequences (8,388,608 tokens) from the validation dataset to perform the validation. According to our experiments, the first 91205 samples from the unshuffled C4 dataset yields 47,186,855 tokens, which is the smallest amount of samples needed to yield 47,185,920 tokens. Thus, we have chosen the first 91205 samples as our validation dataset. + +#### Training data order + +We randomly shuffle the **last 256 of 1024 shards** for the benchmarking area. + +#### Test data order + +We use the first 1024 sequences in the validation dataset for validation. We **do not shuffle** the validation dataset. + +# 4. Model +### Publication/Attribution + +The model largely follows the Llama 3.1 8B [paper](https://arxiv.org/abs/2407.21783). + +### Model details + +| Config | Value | +| :-- | :-- | +| Embedding | RoPE + parameter adjustments | +| # Layers | 32 | +| Attention Type | GQA | +| # Attn Heads | 32 | +| Key/Value Heads | 8 | +| Model Dimension | 4096 | +| FFN Dimension | 14336 | +| Activation | SwiGLU | +| Normalization | RMSNorm | +| Tokenizer | Llama tokenizer | +| Vocab size | 128,000 | +| Context Length | 8192 | + + +#### Saving and restoring a checkpoint + +Large runs might need to span across multiple Slurm jobs, and we need to save and load checkpoints with contexts so that training can resume between jobs. To support this, we have added some environment variables. Please refer to `config.sh` for more details. + +### Optimizer spec + +1. Optimizer type: **AdamW** +2. Warmup steps computed as 10% of the total allocated steps. +3. LR Scheduler's maximum number of steps can be configured in the `config.json`. + +# 5. Quality +### Quality metric + +Validation loss + +### Quality target + +Validation log perplexity = 3.3 + +### Evaluation frequency + +We perform evaluation every **12288** sequences. + +### Evaluation thoroughness + +We evaluate using **1024** sequences from our customized validation dataset. + + +# 6. Other + + + +#### Run model conversion + +Assuming that we have downloaded the HuggingFace checkpoint to a `` directory, we can run [this script](./utils/launch_nemo_convert.sh) (which calls [this python script](./utils/nemo_convert.py)) to perform checkpoint format conversion. After such conversion is done, you should be able to find the converted checkpoint under `` directory, and there should be two subfolders inside this directory - `context` and `weights`. + +```bash +# fill in the built container path here +export CONT_IMAGE_URL="" +# fill in the folder that holds the HF checkpoint here +# under this folder, you should see a lot of safetensors +export SRC_PATH="" +# fill in the destination folder of your choice here +# after conversion is done, you can find context and weights under this path +export DST_PATH="" + +# Extra Slurm-related arguments can be provided here +sbatch launch_nemo_convert.sh +``` + +After the model conversion is done, we can then set `MODEL_CKPT=$DST_PATH` together with `FROM_HF=1` when launching our job, so that we can resume training from the converted HF checkpoint. diff --git a/small_llm_pretraining/nemo/callbacks.py b/small_llm_pretraining/nemo/callbacks.py new file mode 100644 index 000000000..4ba581148 --- /dev/null +++ b/small_llm_pretraining/nemo/callbacks.py @@ -0,0 +1,244 @@ +# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +### MLLogger +from mlperf_logging import mllog +from mlperf_logging.mllog import constants +import torch.distributed as dist + +def is_dist_avail_and_initialized(): + return (dist.is_available() and dist.is_initialized()) + +def get_rank(): + if not is_dist_avail_and_initialized(): + return 0 + return dist.get_rank() + +def barrier(): + if not is_dist_avail_and_initialized(): + return + + dist.barrier() + +class MLLogger: + def __init__(self, filepath="/mlperf-outputs/mlperf_llama31_405b.log", default_stack_offset=2): + self.logger = mllog.get_mllogger() + mllog.config(default_stack_offset=default_stack_offset, filename=filepath) + + def start(self, **kwargs): + if get_rank() == 0: + self.logger.start(**kwargs) + + def end(self, **kwargs): + if get_rank() == 0: + self.logger.end(**kwargs) + + def event(self, **kwargs): + if get_rank() == 0: + self.logger.event(**kwargs) + + def submission_info(self): + self.event(key=constants.SUBMISSION_BENCHMARK, value="llama31_8b") + self.event(key=constants.SUBMISSION_ORG, value="reference_implementation") + self.event(key=constants.SUBMISSION_DIVISION, value=constants.CLOSED) + self.event(key=constants.SUBMISSION_STATUS, value=constants.ONPREM) + self.event(key=constants.SUBMISSION_PLATFORM, value="DGX-H100") + self.event(key=constants.SUBMISSION_POC_NAME, value="Yunzhou Liu") + self.event(key=constants.SUBMISSION_POC_EMAIL, value="yunzhoul@nvidia.com") + +mllogger = MLLogger() + +### Preemptive checkpoint callbacks +import lightning.pytorch as pl +from nemo.utils import logging + +class PreemptiveStop(pl.Callback): + """Preemptively stop training at a given global step. Allows stopping training before reaching + the max steps. Useful for testing checkpoint save and resume. + + Args: + stop_on_step (int): Stop training when trainer.global_step reaches this value. + Checked at the start of every step. + """ + + def __init__(self, stop_on_step: int): + self.stop_on_step = stop_on_step + + def on_train_batch_end( + self, trainer, pl_module, outputs, batch, batch_idx + ) -> None: + if trainer.global_step >= self.stop_on_step: + logging.info(f"Global step {trainer.global_step} >= {self.stop_on_step}, signaling Trainer to stop.") + trainer.should_stop = True + # skip EarlyStopping validation unless val_check_interval met + if trainer.global_step % trainer.val_check_interval != 0: + trainer.limit_val_batches = 0 + + +### Metrics Logger +from pytorch_lightning.loggers import Logger +from pytorch_lightning.utilities import rank_zero_only + +class MetricsLogger(Logger): + def __init__( + self, + init_global_step, global_batch_size, seq_length, + target_log_ppl, + train_loss_key = "reduced_train_loss", + val_loss_key = "val_loss", + train_step_time_in_s = "train_step_timing in s", + train_step_time_atol=7200, + ): + super().__init__() + + self.init_global_step = init_global_step + self.gbs = global_batch_size + self.seq_len = seq_length + + self.target = target_log_ppl + self.train_loss_key = train_loss_key + self.val_loss_key = val_loss_key + self.is_target_reached = False + + self.train_step_time_in_s = train_step_time_in_s + self.train_step_time_atol = train_step_time_atol + + def log_metrics(self, metrics, step): + if self.val_loss_key in metrics: + self.log_validation_loss(metrics, step) + + if self.train_step_time_in_s in metrics: + step_time = metrics[self.train_step_time_in_s] + assert step_time <= self.train_step_time_atol, f"Logged train step time ({step_time}) is slower than tolerable ({self.train_step_time_atol}). " + + def log_validation_loss(self, metrics, step): + consumed_samples = step * self.gbs + + loss = metrics[self.val_loss_key] + + mllogger.event(key=constants.EVAL_ACCURACY, value=loss, metadata={constants.SAMPLES_COUNT: consumed_samples}) + + if not self.is_target_reached and loss <= self.target: + self.is_target_reached = True + + @rank_zero_only + def log_hyperparams(self, params, *args, **kwargs): + pass + + @property + def name(self): + return 'mlperf-metrics' + + @property + def version(self): + return 1 + +### MLPerf callbacks +def compute_consumed_mllog_samples(trainer, init_global_step, global_batch_size, seq_length): + consumed_samples = ( + trainer.global_step * global_batch_size + ) + return int(consumed_samples) # we log the epoch numbers in sequences, not tokens + +class MLPerfCallback(pl.Callback): + def __init__( + self, + global_batch_size, + micro_batch_size, + sequence_length, + init_global_step, + eval_every, + configs={} + ): + mllogger.event(key=constants.CACHE_CLEAR, value=True) + mllogger.start(key=constants.INIT_START) + super().__init__() + + self.init_global_step = init_global_step + self.gbs = global_batch_size + self.mbs = micro_batch_size + self.seq_len = sequence_length + self.eval_every = eval_every + + self.is_target_reached = False + self.status = constants.ABORTED + self.configs = configs + + def consumed_samples(self, trainer): + return compute_consumed_mllog_samples(trainer, self.init_global_step, self.gbs, self.seq_len) + + def set_success_status(self): + self.status = constants.SUCCESS + self.is_target_reached = True + + @rank_zero_only + def on_train_epoch_start(self, trainer, pl_module): + mllogger.start(key=constants.EPOCH_START, metadata={constants.SAMPLES_COUNT: self.consumed_samples(trainer)}) + mllogger.start(key=constants.BLOCK_START, metadata={constants.SAMPLES_COUNT: self.consumed_samples(trainer)}) + + return super().on_train_epoch_start(trainer, pl_module) + + @rank_zero_only + def on_train_epoch_end(self, trainer, pl_module): + mllogger.end(key=constants.EPOCH_STOP, metadata={constants.SAMPLES_COUNT: self.consumed_samples(trainer)}) + return super().on_train_epoch_end(trainer, pl_module) + + def on_train_end(self, trainer, pl_module): + # for every occurrences, run on all ranks to allow sync + barrier() + mllogger.end(key=constants.RUN_STOP, metadata={"status": self.status}) + mllogger.event(key="train_samples", value=self.consumed_samples(trainer)) + return super().on_train_end(trainer, pl_module) + + @rank_zero_only + def log_eval_start(self, trainer, pl_module): + mllogger.end(key=constants.BLOCK_STOP, metadata={constants.SAMPLES_COUNT: self.consumed_samples(trainer)}) + mllogger.start(key=constants.EVAL_START, metadata={constants.SAMPLES_COUNT: self.consumed_samples(trainer)}) + + + def on_validation_start(self, trainer, pl_module): + trainer.val_check_interval = self.eval_every + trainer.val_check_batch = self.eval_every + self.log_eval_start(trainer, pl_module) + + def on_validation_end(self, trainer, pl_module): + mllogger.end(key=constants.EVAL_STOP, metadata={constants.SAMPLES_COUNT: self.consumed_samples(trainer)}) + + for logger in trainer.loggers: + if isinstance(logger, MetricsLogger): + if logger.is_target_reached: + trainer.should_stop = True + self.set_success_status() + + if not trainer.should_stop: + mllogger.start(key=constants.BLOCK_START, metadata={constants.SAMPLES_COUNT: self.consumed_samples(trainer)}) + + return super().on_validation_end(trainer, pl_module) + + @rank_zero_only + def load_state_dict(self, state_dict): + print(f":::MLLOG Weight initialization: {state_dict.keys()}") + return super().load_state_dict(state_dict) + + def on_train_start(self, trainer, pl_module): + # run on all ranks to allow synchronization + barrier() + mllogger.submission_info() + + for key, value in self.configs.items(): + mllogger.event(key=key, value=value) + + mllogger.end(key=constants.INIT_STOP) + mllogger.start(key=constants.RUN_START) diff --git a/small_llm_pretraining/nemo/config_H100_1x8x4_8b.sh b/small_llm_pretraining/nemo/config_H100_1x8x4_8b.sh new file mode 100644 index 000000000..a58eb1e78 --- /dev/null +++ b/small_llm_pretraining/nemo/config_H100_1x8x4_8b.sh @@ -0,0 +1,109 @@ +# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# SSH: username that connects to the remote cluster +export USER="DUMMY" +# SSH: remote cluster URL +export HOST="DUMMY" +# Slurm: account for job submission +export ACCOUNT="DUMMY" +# Slurm: partition for job submission +export PARTITION="DUMMY" +# Slurm: job time limit, defaults to 8 hours +export TIME="08:00:00" +# Slurm: --nodes arguments, default to use 288 nodes +export NNODES=1 +# Slurm: --gpus_per_node and --ntasks_per_node argument, defaults to 8 GPUs per node +export GPUS_PER_NODE=8 +# Slurm: max job retries for transient job failures, defaults to retry 3 times +export MAX_RETRIES=1 + +# Folder mapping: +# Output directory that holds logs, any path that you like. +export JOB_DIR="/workspace/code/logs" +# Image / container path, either local cache file or remote URL +export IMAGE="DUMMY" +# Dataset: C4 dataset location that contains the dataset after preprocessing +# export ORIGINAL_C4_PATH="/data/data/C4" + +# This corresponds to the PREPROCESSED_PATH in README section 3's dataset download part +export PREPROCESSED_PATH="/data/llama3_8b/data/C4_processed" +export MERGED_C4_PATH="/data/llama3_8b/data/C4_merged" +# Dataset: Numpy index working directory, contains shuffled dataset +# This path must be able to hold >400GB data +export TMP_NPY_INDEX="/data/npy_indices" +# Dataset: Tokenizer path +# This corresponds to the TOKENIZER_PATH in README section 3's tokenizer download part +export TOKENIZER_PATH="/data/llama3_8b/model/Llama-3.1-8B" +# export TOKENIZER_PATH="/data/llama3_405b_ref/tokenizer" + +# Model: checkpoint and tokenizer path +# This is the checkpoint that we want to start with. +# Each checkpoint should be a folder containing two sub-folders: context and weights. +# And we need to pass this folder's path (the folder containing context and weights) here. +export MODEL_CKPT="/data/llama3_8b/model/Llama-3.1-8B_nemo" +# export MODEL_CKPT="None" +# Model: Continual checkpoint directory to write and resume +# This is the directory to hold all intermediate checkpoints. +# Once a run is complete and we specify to save checkpoints, +# we should see a checkpoint written in this folder +# with name `checkpoint-par-x-y-steps` +# Inside this directory, there should be a `checkpoint` directory that holds context and weights +# which is the "actual checkpoint". +# Notice that this path must be able to hold at least 5.2TB data since each checkpoint is 5.2TB. +export CONTINUAL_CKPT="/data/model/saved_ckpts" +# Model: Whether we want to restore from MODEL_CKPT path. If 0, then we are not restoring. +export USE_CKPT=0 +# Model: Whether we are resuming from a NeMo-formatted HuggingFace checkpoint (weights only). +# If set to 1, then checkpoint resuming code will not try to load the optimizer states. +export FROM_HF=1 +# Model: Whether we want to save a checkpoint. Must be 1 if NPAR > 1. If 1, then we save a checkpoint at the end. +export SAVE_CKPT=0 + + + +# Training Configs: +# Model: size, to choose from 8b, 70b, 405b +export SIZE="8b" +# Dataloader: Global batch size +export GBS=32 +# Dataloader: Micro batch size +export MBS=1 +export MAX_LR="5e-4" +# Dataloader: Max run N batches, optional +# If an empty string is provided (""), then the training will continue until time limit +# If we want to save a checkpoint, then this value must be set +# export MAX_STEPS=1200000 # Fixed max_steps=1200000 in pretrain_llama31.py +export WARMUP_STEPS=512 # 16384 // GBS +export EVAL_EVERY=12288 +export START_EVAL_AT=0 + +export TENSOR_PARALLEL_SIZE=4 + +# Experiment: starting steps +# This is the starting "offset" step from the checkpoint. +# For instance, if you are resuming from a checkpoint folder `checkpoint-par-0-20-steps/checkpoint`, +# which means that the model is trained for 20 steps to generate the checkpoint, +# then the value 20 is needed here. +export START_STEPS="0" +# Experiment manager: Number of experiments to launch +export NEXP=1 +# Experiment manager: how many consecutive jobs we want for each experiment +export NPAR=1 +# Experiment manager: provides seeds to the launched experiments, use space as delimiter, such as "1234 1235 1236" +# The training script will discard all excessive seeds, and generate seeds if given seeds < NEXP. +# To preserve randomness, we recommend not to set this value so that each time seeds can be randomly generated. + + +export DGXSYSTEM=$(basename $(readlink -f ${BASH_SOURCE[0]}) | sed 's/^config_//' | sed 's/\.sh$//' ) diff --git a/small_llm_pretraining/nemo/config_H200_1x8x1_8b.sh b/small_llm_pretraining/nemo/config_H200_1x8x1_8b.sh new file mode 100644 index 000000000..a9689f1c9 --- /dev/null +++ b/small_llm_pretraining/nemo/config_H200_1x8x1_8b.sh @@ -0,0 +1,109 @@ +# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# SSH: username that connects to the remote cluster +export USER="DUMMY" +# SSH: remote cluster URL +export HOST="DUMMY" +# Slurm: account for job submission +export ACCOUNT="DUMMY" +# Slurm: partition for job submission +export PARTITION="DUMMY" +# Slurm: job time limit, defaults to 8 hours +export TIME="08:00:00" +# Slurm: --nodes arguments, default to use 288 nodes +export NNODES=1 +# Slurm: --gpus_per_node and --ntasks_per_node argument, defaults to 8 GPUs per node +export GPUS_PER_NODE=8 +# Slurm: max job retries for transient job failures, defaults to retry 3 times +export MAX_RETRIES=1 + +# Folder mapping: +# Output directory that holds logs, any path that you like. +export JOB_DIR="/workspace/code/logs" +# Image / container path, either local cache file or remote URL +export IMAGE="DUMMY" +# Dataset: C4 dataset location that contains the dataset after preprocessing +# export ORIGINAL_C4_PATH="/data/data/C4" + +# This corresponds to the PREPROCESSED_PATH in README section 3's dataset download part +export PREPROCESSED_PATH="/data/llama3_8b/data/C4_processed" +export MERGED_C4_PATH="/data/llama3_8b/data/C4_merged" +# Dataset: Numpy index working directory, contains shuffled dataset +# This path must be able to hold >400GB data +export TMP_NPY_INDEX="/data/npy_indices" +# Dataset: Tokenizer path +# This corresponds to the TOKENIZER_PATH in README section 3's tokenizer download part +export TOKENIZER_PATH="/data/llama3_8b/model/Llama-3.1-8B" +# export TOKENIZER_PATH="/data/llama3_405b_ref/tokenizer" + +# Model: checkpoint and tokenizer path +# This is the checkpoint that we want to start with. +# Each checkpoint should be a folder containing two sub-folders: context and weights. +# And we need to pass this folder's path (the folder containing context and weights) here. +export MODEL_CKPT="/data/llama3_8b/model/Llama-3.1-8B_nemo" +# export MODEL_CKPT="None" +# Model: Continual checkpoint directory to write and resume +# This is the directory to hold all intermediate checkpoints. +# Once a run is complete and we specify to save checkpoints, +# we should see a checkpoint written in this folder +# with name `checkpoint-par-x-y-steps` +# Inside this directory, there should be a `checkpoint` directory that holds context and weights +# which is the "actual checkpoint". +# Notice that this path must be able to hold at least 5.2TB data since each checkpoint is 5.2TB. +export CONTINUAL_CKPT="/data/model/saved_ckpts" +# Model: Whether we want to restore from MODEL_CKPT path. If 0, then we are not restoring. +export USE_CKPT=0 +# Model: Whether we are resuming from a NeMo-formatted HuggingFace checkpoint (weights only). +# If set to 1, then checkpoint resuming code will not try to load the optimizer states. +export FROM_HF=1 +# Model: Whether we want to save a checkpoint. Must be 1 if NPAR > 1. If 1, then we save a checkpoint at the end. +export SAVE_CKPT=0 + + + +# Training Configs: +# Model: size, to choose from 8b, 70b, 405b +export SIZE="8b" +# Dataloader: Global batch size +export GBS=32 +# Dataloader: Micro batch size +export MBS=2 +export MAX_LR="5e-4" +# Dataloader: Max run N batches, optional +# If an empty string is provided (""), then the training will continue until time limit +# If we want to save a checkpoint, then this value must be set +# export MAX_STEPS=1200000 # Fixed max_steps=1200000 in pretrain_llama31.py +export WARMUP_STEPS=512 # 16384 // GBS +export EVAL_EVERY=12288 +export START_EVAL_AT=0 + +export TENSOR_PARALLEL_SIZE=1 + +# Experiment: starting steps +# This is the starting "offset" step from the checkpoint. +# For instance, if you are resuming from a checkpoint folder `checkpoint-par-0-20-steps/checkpoint`, +# which means that the model is trained for 20 steps to generate the checkpoint, +# then the value 20 is needed here. +export START_STEPS="0" +# Experiment manager: Number of experiments to launch +export NEXP=1 +# Experiment manager: how many consecutive jobs we want for each experiment +export NPAR=1 +# Experiment manager: provides seeds to the launched experiments, use space as delimiter, such as "1234 1235 1236" +# The training script will discard all excessive seeds, and generate seeds if given seeds < NEXP. +# To preserve randomness, we recommend not to set this value so that each time seeds can be randomly generated. + + +export DGXSYSTEM=$(basename $(readlink -f ${BASH_SOURCE[0]}) | sed 's/^config_//' | sed 's/\.sh$//' ) diff --git a/small_llm_pretraining/nemo/config_MI325X_1x8x1_8b.sh b/small_llm_pretraining/nemo/config_MI325X_1x8x1_8b.sh new file mode 100644 index 000000000..5f6034613 --- /dev/null +++ b/small_llm_pretraining/nemo/config_MI325X_1x8x1_8b.sh @@ -0,0 +1,109 @@ +# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# SSH: username that connects to the remote cluster +export USER="DUMMY" +# SSH: remote cluster URL +export HOST="DUMMY" +# Slurm: account for job submission +export ACCOUNT="DUMMY" +# Slurm: partition for job submission +export PARTITION="DUMMY" +# Slurm: job time limit, defaults to 8 hours +export TIME="08:00:00" +# Slurm: --nodes arguments, default to use 288 nodes +export NNODES=1 +# Slurm: --gpus_per_node and --ntasks_per_node argument, defaults to 8 GPUs per node +export GPUS_PER_NODE=8 +# Slurm: max job retries for transient job failures, defaults to retry 3 times +export MAX_RETRIES=1 + +# Folder mapping: +# Output directory that holds logs, any path that you like. +export JOB_DIR="/workspace/code/logs" +# Image / container path, either local cache file or remote URL +export IMAGE="DUMMY" +# Dataset: C4 dataset location that contains the dataset after preprocessing +# export ORIGINAL_C4_PATH="/data/data/C4" + +# This corresponds to the PREPROCESSED_PATH in README section 3's dataset download part +export PREPROCESSED_PATH="/data/llama31_8b/data/C4_processed/" +export MERGED_C4_PATH="/data/llama31_8b/data/C4_merged" +# Dataset: Numpy index working directory, contains shuffled dataset +# This path must be able to hold >400GB data +export TMP_NPY_INDEX="/data/npy_indices" +# Dataset: Tokenizer path +# This corresponds to the TOKENIZER_PATH in README section 3's tokenizer download part +export TOKENIZER_PATH="/data/llama31_8b/model/Llama-3.1-8B-ref/" +# export TOKENIZER_PATH="/data/llama3_405b_ref/tokenizer" + +# Model: checkpoint and tokenizer path +# This is the checkpoint that we want to start with. +# Each checkpoint should be a folder containing two sub-folders: context and weights. +# And we need to pass this folder's path (the folder containing context and weights) here. +export MODEL_CKPT="/data/llama31_8b/model/Llama-3.1-8B-ref/" +# export MODEL_CKPT="None" +# Model: Continual checkpoint directory to write and resume +# This is the directory to hold all intermediate checkpoints. +# Once a run is complete and we specify to save checkpoints, +# we should see a checkpoint written in this folder +# with name `checkpoint-par-x-y-steps` +# Inside this directory, there should be a `checkpoint` directory that holds context and weights +# which is the "actual checkpoint". +# Notice that this path must be able to hold at least 5.2TB data since each checkpoint is 5.2TB. +export CONTINUAL_CKPT="/data/model/saved_ckpts" +# Model: Whether we want to restore from MODEL_CKPT path. If 0, then we are not restoring. +export USE_CKPT=0 +# Model: Whether we are resuming from a NeMo-formatted HuggingFace checkpoint (weights only). +# If set to 1, then checkpoint resuming code will not try to load the optimizer states. +export FROM_HF=1 +# Model: Whether we want to save a checkpoint. Must be 1 if NPAR > 1. If 1, then we save a checkpoint at the end. +export SAVE_CKPT=0 + + + +# Training Configs: +# Model: size, to choose from 8b, 70b, 405b +export SIZE="8b" +# Dataloader: Global batch size +export GBS=128 +# Dataloader: Micro batch size +export MBS=4 +export MAX_LR="5e-4" +# Dataloader: Max run N batches, optional +# If an empty string is provided (""), then the training will continue until time limit +# If we want to save a checkpoint, then this value must be set +# export MAX_STEPS=1200000 # Fixed max_steps=1200000 in pretrain_llama31.py +export WARMUP_STEPS=512 # 16384 // GBS +export EVAL_EVERY=12288 +export START_EVAL_AT=0 + +export TENSOR_PARALLEL_SIZE=1 +# Experiment: starting steps +# This is the starting "offset" step from the checkpoint. +# For instance, if you are resuming from a checkpoint folder `checkpoint-par-0-20-steps/checkpoint`, +# which means that the model is trained for 20 steps to generate the checkpoint, +# then the value 20 is needed here. +export START_STEPS="0" +# Experiment manager: Number of experiments to launch +export NEXP=1 +# Experiment manager: how many consecutive jobs we want for each experiment +export NPAR=1 +# Experiment manager: provides seeds to the launched experiments, use space as delimiter, such as "1234 1235 1236" +# The training script will discard all excessive seeds, and generate seeds if given seeds < NEXP. +# To preserve randomness, we recommend not to set this value so that each time seeds can be randomly generated. +# export SEEDS=7963 +# export SEEDS=1234 + +export DGXSYSTEM=$(basename $(readlink -f ${BASH_SOURCE[0]}) | sed 's/^config_//' | sed 's/\.sh$//' ) diff --git a/small_llm_pretraining/nemo/debug.py b/small_llm_pretraining/nemo/debug.py new file mode 100644 index 000000000..9a304244a --- /dev/null +++ b/small_llm_pretraining/nemo/debug.py @@ -0,0 +1,10 @@ +from nemo.collections.common.tokenizers import AutoTokenizer +from pprint import pprint +tokenizer = AutoTokenizer(pretrained_model_name="/data/llama3_8b_ref/model/Llama-3.1-8B") +print (f'tokenizer: {tokenizer}') +print(dir(tokenizer)) + + +print (tokenizer.vocab_size) +# pprint(vars(tokenizer)) + diff --git a/small_llm_pretraining/nemo/dev/run_docker.sh b/small_llm_pretraining/nemo/dev/run_docker.sh new file mode 100644 index 000000000..26a5a8e70 --- /dev/null +++ b/small_llm_pretraining/nemo/dev/run_docker.sh @@ -0,0 +1,12 @@ +# Change directory to the model directory +SCRIPT_DIR=$(dirname "$(readlink -f "$0")") +cd $SCRIPT_DIR/.. + +docker run -it --rm \ + --net=host --uts=host \ + --ipc=host --device /dev/dri --device /dev/kfd \ + --security-opt=seccomp=unconfined \ + --volume=/data/training:/data \ + --volume $(pwd):/workspace/code/ \ + --volume=/data/training/llama3_8b/outputs:/outputs \ + # --name llama-training-`whoami` rocm/mlperf:llama31_8b_training_5.1_gfx942_v2 diff --git a/small_llm_pretraining/nemo/dev/run_llama31.sh b/small_llm_pretraining/nemo/dev/run_llama31.sh new file mode 100755 index 000000000..66ab35223 --- /dev/null +++ b/small_llm_pretraining/nemo/dev/run_llama31.sh @@ -0,0 +1,154 @@ +#!/bin/bash + +# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +set -e + +#git config --global --add safe.directory /workspace/llama31 + +# Vars without defaults +# Slurm settings +: "${USER:?USER not set}" +: "${HOST:?HOST not set}" +: "${ACCOUNT:?ACCOUNT not set}" +: "${PARTITION:?PARTITION not set}" +: "${REMOTE:=0}" + +# Job settings +: "${JOB_DIR:?JOB_DIR not set}" +: "${IMAGE:?IMAGE not set}" + +# Dataset settings +: "${PREPROCESSED_PATH:?PREPROCESSED_PATH not set}" +: "${TOKENIZER_PATH:?TOKENIZER_PATH not set}" + +# Model settings +: "${MODEL_CKPT:?MODEL_CKPT not set}" +: "${USE_CKPT:?USE_CKPT not set}" +: "${FROM_HF:?FROM_HF not set}" +: "${CONTINUAL_CKPT:?CONTINUAL_CKPT not set}" + +# Vars with defaults +# Slurm settings +: "${TIME:="04:00:00"}" +: "${NNODES:=1}" +: "${GPUS_PER_NODE:=8}" +: "${DEPENDENCIES:=""}" + +# Job settings +: "${NEMO_DIR:=""}" # Provide customized NeMo path here +: "${NEMO_RUN_DIR:=""}" # Provide customized NeMo-Run path here +: "${TMP_NPY_INDEX:=""}" # Provide temporary NNumpy Index saving directory +: "${MAX_RETRIES:=0}" + +# Model settings +: "${SIZE:="8b"}" +: "${GBS:=4}" +: "${MBS:=1}" +: "${START_STEPS:=0}" + +# Dataloader settings +: "${MAX_STEPS:=""}" + +# Experiment settings +: "${SEEDS:=""}" +IFS=" " read -ra seeds <<< $SEEDS +: "${NEXP:=1}" +: "${NPAR:=1}" +: "${SAVE_CKPT:=0}" +: "${TAG:=""}" +: "${TARGET:="3.3"}" +: "${STEP_TIME_ATOL:="18000"}" # maximum tolerable step time, setting to 2hr by default + +# Run + +MOUNTS="${JOB_DIR}:/output,${JOB_DIR}:/mlperf-outputs,${PREPROCESSED_PATH}:/preproc_data,${MODEL_CKPT}:/checkpoint,${TOKENIZER_PATH}:/tokenizer,${CONTINUAL_CKPT}:/continual" + +CKPT_OPTION="" + +CMD_SUFFIX="" + +if [ $USE_CKPT -gt 0 ]; then + CMD_SUFFIX="${CMD_SUFFIX} --use_ckpt" + if [ $FROM_HF -gt 0 ]; then + CMD_SUFFIX="${CMD_SUFFIX} --resume_from_hf" + fi +fi + +if [ $SAVE_CKPT -gt 0 ]; then + CMD_SUFFIX="${CMD_SUFFIX} --save_ckpt" +fi + +if [ ! $NEMO_DIR = "" ]; then + MOUNTS="${MOUNTS},${NEMO_DIR}:/opt/NeMo" +fi + +if [ ! $NEMO_RUN_DIR = "" ]; then + MOUNTS="${MOUNTS},${NEMO_RUN_DIR}:/opt/NeMo-Run" +fi + +if [ ! $TMP_NPY_INDEX = "" ]; then + MOUNTS="${MOUNTS},${TMP_NPY_INDEX}:/npy_index" +fi + +if [ ! $DEPENDENCIES = "" ]; then + CMD_SUFFIX="${CMD_SUFFIX} --dependencies ${DEPENDENCIES}" +fi + +if [ ! $MAX_STEPS = "" ]; then + CMD_SUFFIX="${CMD_SUFFIX} --max_steps ${MAX_STEPS}" +fi + +if [ ! $TAG = "" ]; then + CMD_SUFFIX="${CMD_SUFFIX} --tag ${TAG}" +fi + +if [ $REMOTE -gt 0 ]; then + CMD_SUFFIX="${CMD_SUFFIX} --run_slurm" +fi + +if [ $TENSOR_PARALLEL_SIZE -gt 0 ]; then + CMD_SUFFIX="${CMD_SUFFIX} --tensor_parallel_size ${TENSOR_PARALLEL_SIZE}" +fi + +# Allows MLLogger objects to be constructed locally +if [ ! -d /mlperf-outputs ]; then mkdir /mlperf-outputs; fi + +set -x + +python3 pretrain_llama31.py \ +--user $USER --host $HOST \ +--job_dir $JOB_DIR \ +--account $ACCOUNT --partition $PARTITION \ +--nodes $NNODES --gpus_per_node $GPUS_PER_NODE \ +--time $TIME \ +--max_retries $MAX_RETRIES \ +--mounts $MOUNTS \ +--image $IMAGE \ +--size $SIZE \ +--gbs $GBS --mbs $MBS \ +--seeds ${seeds[@]} \ +--num_exps $NEXP \ +--num_pars $NPAR \ +--initial_ckpt_path $MODEL_CKPT \ +--continual_ckpt_path $CONTINUAL_CKPT \ +--tokenizer_path $TOKENIZER_PATH \ +--target_log_ppl $TARGET \ +--step_time_atol $STEP_TIME_ATOL \ +--ckpt_start_step $START_STEPS \ +--warmup_steps $WARMUP_STEPS \ +--eval_every $EVAL_EVERY \ +--start_eval_at $START_EVAL_AT \ +$CMD_SUFFIX diff --git a/small_llm_pretraining/nemo/dev/run_with_docker_sub.sh b/small_llm_pretraining/nemo/dev/run_with_docker_sub.sh new file mode 100755 index 000000000..04cead521 --- /dev/null +++ b/small_llm_pretraining/nemo/dev/run_with_docker_sub.sh @@ -0,0 +1,110 @@ +#!/bin/bash + +# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +set -euxo pipefail + +# Change directory to the model directory +SCRIPT_DIR=$(dirname "$(readlink -f "$0")") +cd $SCRIPT_DIR/.. + +# Vars without defaults +: "${DGXSYSTEM:?DGXSYSTEM not set}" +: "${CONT:?CONT not set}" +: "${DATADIR:?DATADIR not set}" +: "${MODEL:?MODEL not set}" +: "${TOKENIZER:?TOKENIZER not set}" + +# Vars with defaults +: "${NEXP:=1}" +: "${DATESTAMP:=$(date +'%y%m%d%H%M%S%N')}" +: "${CLEAR_CACHES:=1}" +: "${CHECK_COMPLIANCE:=1}" +: "${MLPERF_RULESET:=5.0.0}" +: "${LOGDIR:=./results}" +: "${DEPENDENCIES:=./dependencies}" +: "${CONT_NAME:=dev-${CUSTOM_TAG}}" +: "${LOG_FREQ:=0}" + +# Other vars +readonly _config_file="./config_${DGXSYSTEM}.sh" +readonly _logfile_base="${LOGDIR}/${DATESTAMP}" +readonly _cont_name="${CONT_NAME}" +_cont_mounts=("--volume=${DATADIR}:/data/data" "--volume=${MODEL}:/data/model/" "--volume=${TOKENIZER}:/data/tokenizer/" "--volume=$(pwd):/workspace/code" "--volume=$(pwd)/../../AMD:/workspace/AMD" "--volume=$(pwd)/../../utilities:/workspace/utilities" "--volume=${LOGDIR}:/results") + + +# Setup directories +mkdir -p "${LOGDIR}" +mkdir -p "${LOGDIR}/artifacts/" + +# Get list of envvars to pass to docker +mapfile -t _config_env < <(env -i bash -c ". ${_config_file} && compgen -e" | grep -E -v '^(PWD|SHLVL)') +_config_env+=(DATADIR) +_config_env+=(MODEL) +_config_env+=(DGXSYSTEM) +_config_env+=(LOGDIR) + +echo "TEST" +echo ${_config_env[@]} +mapfile -t _config_env < <(for v in "${_config_env[@]}"; do echo "--env=$v"; done) + +# Cleanup container +cleanup_docker() { + if docker ps -a --format '{{.Names}}' | grep -q "^${_cont_name}$"; then + docker container rm -f "${_cont_name}" || true + else + echo "Container ${_cont_name} does not exist. Skipping removal." + fi +} +cleanup_docker +trap 'set -eux; cleanup_docker' EXIT + +docker run --rm --init --detach \ + --net=host --uts=host \ + --ipc=host --device /dev/dri --device /dev/kfd \ + --security-opt=seccomp=unconfined \ + --name="${_cont_name}" "${_cont_mounts[@]}" \ + -e IMAGE_NAME="${CONT}" \ + "${CONT}" sleep infinity + +# Make sure container has time to finish initialization +sleep 5 +# bash runtime_tunables.sh +docker exec "${_cont_name}" true + +# Run experiments +for _experiment_index in $(seq 1 "${NEXP}"); do + ( + echo "Beginning trial ${_experiment_index} of ${NEXP}" + if [[ $CLEAR_CACHES == 1 ]]; then + bash -c "echo -n 'Clearing cache on ' && hostname && sync && sudo /sbin/sysctl vm.drop_caches=3" + fi + _config_env+=(--env=SEED=$RANDOM) # Reset random seed + echo 'launching experiment using:' ${_config_env[@]} ${_cont_name} ./dev/run_llama31.sh + docker exec ${_config_env[@]} --env=HYDRA_FULL_ERROR ${_cont_name} ./dev/run_llama31.sh + ) | tee "${_logfile_base}_${_experiment_index}.log" + + if [ "${CHECK_COMPLIANCE}" -eq 1 ]; then + docker exec "${_config_env[@]}" "${_cont_name}" \ + python3 -m mlperf_logging.compliance_checker --usage training \ + --ruleset "${MLPERF_RULESET}" \ + --log_output "/results/compliance_${DATESTAMP}.out" \ + "/results/${DATESTAMP}_${_experiment_index}.log" \ + || true + fi + +done + +echo "Number of experiments $NEXP" diff --git a/small_llm_pretraining/nemo/mcore.patch b/small_llm_pretraining/nemo/mcore.patch new file mode 100644 index 000000000..5a9639e64 --- /dev/null +++ b/small_llm_pretraining/nemo/mcore.patch @@ -0,0 +1,90 @@ +diff --git a/megatron/core/datasets/gpt_dataset.py b/megatron/core/datasets/gpt_dataset.py +index 2eb7702b..d1f0b9a9 100644 +--- a/megatron/core/datasets/gpt_dataset.py ++++ b/megatron/core/datasets/gpt_dataset.py +@@ -407,9 +407,10 @@ class GPTDataset(MegatronDataset): + + numpy_random_state = numpy.random.RandomState(self.config.random_seed) + ++ shuffle = self.index_split == Split.train + # Build the document index + document_index = _build_document_index( +- self.indices, num_epochs, numpy_random_state, separate_final_epoch ++ self.indices, num_epochs, numpy_random_state, separate_final_epoch, shuffle + ) + + drop_last_partial_sequence = True +@@ -450,11 +451,11 @@ class GPTDataset(MegatronDataset): + # Build the shuffle index + if separate_final_epoch: + shuffle_index = _build_shuffle_index( +- num_samples_sans_final_epoch, sample_index.shape[0] - 1, numpy_random_state ++ num_samples_sans_final_epoch, sample_index.shape[0] - 1, numpy_random_state, shuffle + ) + else: + shuffle_index = _build_shuffle_index( +- sample_index.shape[0] - 1, sample_index.shape[0] - 1, numpy_random_state ++ sample_index.shape[0] - 1, sample_index.shape[0] - 1, numpy_random_state, shuffle + ) + + if path_to_cache: +@@ -558,6 +559,7 @@ def _build_document_index( + num_epochs: int, + numpy_random_state: numpy.random.RandomState, + separate_final_epoch: bool, ++ shuffle: bool = True, + ) -> numpy.ndarray: + """Build an array with length = num epochs * num documents + +@@ -578,7 +580,8 @@ def _build_document_index( + document_index[:] = documents + document_index = document_index.reshape(-1) + document_index = document_index.astype(numpy.int32) +- numpy_random_state.shuffle(document_index) ++ if shuffle: ++ numpy_random_state.shuffle(document_index) + return document_index + + doc_idx_first = _build_document_index(documents, num_epochs - 1, numpy_random_state, False) +@@ -587,7 +590,8 @@ def _build_document_index( + + + def _build_shuffle_index( +- num_samples: int, total_size: int, numpy_random_state: numpy.random.RandomState ++ num_samples: int, total_size: int, numpy_random_state: numpy.random.RandomState, ++ shuffle: bool = True + ) -> numpy.ndarray: + """Build the range [0, size) and shuffle + +@@ -607,12 +611,16 @@ def _build_shuffle_index( + dtype_ = numpy.int64 + + shuffle_idx_first = numpy.arange(start=0, stop=num_samples, step=1, dtype=dtype_) +- numpy_random_state.shuffle(shuffle_idx_first) ++ ++ if shuffle: ++ numpy_random_state.shuffle(shuffle_idx_first) + if num_samples == total_size: + return shuffle_idx_first + + shuffle_idx_last = numpy.arange(start=num_samples, stop=total_size, step=1, dtype=dtype_) +- numpy_random_state.shuffle(shuffle_idx_last) ++ ++ if shuffle: ++ numpy_random_state.shuffle(shuffle_idx_last) + + return numpy.concatenate((shuffle_idx_first, shuffle_idx_last)) + +diff --git a/megatron/core/transformer/moe/moe_utils.py b/megatron/core/transformer/moe/moe_utils.py +index 0c1504d4..71d29629 100644 +--- a/megatron/core/transformer/moe/moe_utils.py ++++ b/megatron/core/transformer/moe/moe_utils.py +@@ -264,6 +264,7 @@ def topk_softmax_with_capacity( + # Pre softmax + scores = torch.softmax(logits, dim=-1, dtype=torch.float32).type_as(logits) + probs, top_indices = torch.topk(scores, k=topk, dim=1) ++ probs /= probs.sum(dim=-1, keepdim=True) + else: + # Post softmax + if topk == 1: + diff --git a/small_llm_pretraining/nemo/patches/nemo_v2_1_0.patch b/small_llm_pretraining/nemo/patches/nemo_v2_1_0.patch new file mode 100644 index 000000000..590256d55 --- /dev/null +++ b/small_llm_pretraining/nemo/patches/nemo_v2_1_0.patch @@ -0,0 +1,64 @@ +diff --git a/nemo/collections/multimodal/modules/stable_diffusion/attention.py b/nemo/collections/multimodal/modules/stable_diffusion/attention.py +index 646540e88..fad8e4e0a 100644 +--- a/nemo/collections/multimodal/modules/stable_diffusion/attention.py ++++ b/nemo/collections/multimodal/modules/stable_diffusion/attention.py +@@ -143,7 +143,13 @@ def zero_module(module): + + + def Normalize(in_channels, num_groups=32, act=""): +- return GroupNorm(num_groups=num_groups, num_channels=in_channels, eps=1e-6, affine=True, act=act) ++ try: ++ from apex.contrib.group_norm import GroupNorm as GroupNorm ++ return GroupNorm(num_groups=num_groups, num_channels=in_channels, eps=1e-6, affine=True, act=act) ++ except ImportError: ++ print("Using torch.nn.GroupNorm. Hip/Cuda kernel could not be imported from Apex") ++ import torch.nn.GroupNorm as GroupNorm ++ return GroupNorm(num_groups=num_groups, num_channels=in_channels, eps=1e-6, affine=True) + + + class LinearAttention(nn.Module): +@@ -595,4 +601,4 @@ class SpatialTransformer(nn.Module): + x = x.transpose(1, 2).view(b, c, h, w) # b (h w) c -> b c h w + if not self.use_linear: + x = self.proj_out(x) +- return x_in + x ++ return x_in + x +\ No newline at end of file +diff --git a/nemo/collections/multimodal/modules/stable_diffusion/diffusionmodules/util.py b/nemo/collections/multimodal/modules/stable_diffusion/diffusionmodules/util.py +index 69700a436..bfd21ab7e 100644 +--- a/nemo/collections/multimodal/modules/stable_diffusion/diffusionmodules/util.py ++++ b/nemo/collections/multimodal/modules/stable_diffusion/diffusionmodules/util.py +@@ -257,7 +257,13 @@ def mean_flat(tensor): + + + def normalization(in_channels, act="", gn_groups=32): +- return GroupNorm(num_groups=gn_groups, num_channels=in_channels, eps=1e-5, affine=True, act=act) ++ try: ++ from apex.contrib.group_norm import GroupNorm as GroupNorm ++ return GroupNorm(num_groups=num_groups, num_channels=in_channels, eps=1e-6, affine=True, act=act) ++ except ImportError: ++ print("Using torch.nn.GroupNorm. Hip/Cuda kernel could not be imported from Apex") ++ import torch.nn.GroupNorm as GroupNorm ++ return GroupNorm(num_groups=num_groups, num_channels=in_channels, eps=1e-6, affine=True) + + + # PyTorch 1.7 has SiLU, but we support PyTorch 1.5. +@@ -361,4 +367,4 @@ def exists(x): + def default(val, d): + if exists(val): + return val +- return d() if isfunction(d) else d ++ return d() if isfunction(d) else d +\ No newline at end of file +diff --git a/requirements/requirements_nlp.txt b/requirements/requirements_nlp.txt +index 6a86dacbf..a6f60380e 100644 +--- a/requirements/requirements_nlp.txt ++++ b/requirements/requirements_nlp.txt +@@ -11,7 +11,6 @@ jieba + mamba-ssm==2.2.2; sys_platform == 'linux' + markdown2 + matplotlib>=3.3.2 +-#megatron_core>0.6.0 # add back once mcore on pypi is compatible again + nltk>=3.6.5 + numpy<2 # tensorstore has an implicit compiled dependency on numpy<2 + opencc diff --git a/small_llm_pretraining/nemo/pretrain_llama31.py b/small_llm_pretraining/nemo/pretrain_llama31.py new file mode 100644 index 000000000..19ee5a817 --- /dev/null +++ b/small_llm_pretraining/nemo/pretrain_llama31.py @@ -0,0 +1,609 @@ +# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import os +import math +import argparse +from typing import Optional + +import torch +import wandb +from lightning.pytorch.loggers import WandbLogger + +from nemo.collections import llm +from nemo.collections.common.tokenizers import AutoTokenizer +from nemo import lightning as nl +from nemo.collections.llm.recipes.optim.adam import distributed_fused_adam_with_cosine_annealing +import nemo_run as run +from nemo.lightning.run import plugins +from nemo.collections.llm.gpt.data import build_pretraining_datamodule + +from callbacks import PreemptiveStop, MLPerfCallback, MetricsLogger + + +def local_executor( + custom_env_vars: Optional[dict[str, str]] = None, + devices: int = 8, + retries: int = 0, +) -> run.LocalExecutor: + env_vars = { + "TRANSFORMERS_OFFLINE": "1", + "TORCH_NCCL_AVOID_RECORD_STREAMS": "1", + "NCCL_NVLS_ENABLE": "0", + "TOKENIZERS_PARALLELISM": "false", + "NCCL_MIN_P2P_NCHANNELS" : "32", + "NCCL_MIN_CTAS" : "32", + "NCCL_NCHANNELS_PER_NET_PEER" : "32", + "CUBLAS_FORCE_XMMA_KERNEL_INIT" : "DEVICE", + "NVTE_RS_STRIDED_ATOMIC" : "2", + "NVTE_FP8_DPA_BWD" : "1", + "NVTE_FUSED_ATTN" : "1", + "NVTE_FUSED_ATTN_CK" : "1", + "NVTE_FUSED_ATTN_AOTRITON" : "1", + "NVTE_DEBUG" : "0", + "NVTE_DEBUG_LEVEL" : "0", + "NVTE_USE_HIPBLASLT" : "1", + "NVTE_USE_CAST_TRANSPOSE_TRITON" : "0", + "NVTE_USE_OPTIMIZED_HIPIFIED_CAST_TRANSPOSE" : "1", + "USE_TE_SWIGLU" : "1", + "NVTE_CK_USES_BWD_V3" : "1", # enable dqdkdv bwd kernel + "NVTE_CK_V3_BF16_CVT" : "1", # Use Round to away from ZERO for numerical stability + "CK_FUSED_ATTN_LOG_CONFIG" : "0", # Diable logging for CK fused attn. Enabled for debugging only + "NVTE_CK_IS_V3_ATOMIC_FP32" : "0", # 16bit atomics + "NVTE_CK_HOW_V3_BF16_CVT" : "2", + "NVTE_CK_USES_FWD_V3" : "1", + "CUDNN_FRONTEND_ATTN_DP_WORKSPACE_LIMIT" : "0", + "CUDA_DEVICE_MAX_CONNECTIONS" : "1", + "FUSED_SOFTMAX" : "0", + "RMSNORM_CAST" : "0", + "PT_TENSOR_VALIDATION" : "0", + "USE_HIPBLASLT" : "1", + "TORCH_BLAS_PREFER_HIPBLASLT" : "1", + "NVTE_USE_RMSNORM_TRITON" : "1", + "ENABLE_TRANSPOSE_CACHE" : "0", + "NVTE_UNFUSED_FP8_UPDATE": "1", + # "NVTE_DP_AMAX_REDUCE_INTERVAL": "0", + # "NVTE_ASYNC_AMAX_REDUCTION": "1", + } + if custom_env_vars: + env_vars |= custom_env_vars + + executor = run.LocalExecutor() + executor.launcher = 'torchrun' + executor.env_vars = env_vars + executor.retries = retries + executor.nodes = 1 + executor.ntasks_per_node = devices + + return executor + +def slurm_executor( + user: str, + host: str, + remote_job_dir: str, + account: str, + partition: str, + nodes: int, + devices: int, + time: str = "01:00:00", + custom_mounts: Optional[list[str]] = None, + custom_env_vars: Optional[dict[str, str]] = None, + container_image: str = "nvcr.io/nvidia/nemo:dev", + dependencies: list[str] = [], + retries: int = 0, +) -> run.SlurmExecutor: + if not (user and host and remote_job_dir and account and partition and nodes and devices): + raise RuntimeError( + "Please set user, host, remote_job_dir, account, partition, nodes and devices args for using this function." + ) + + mounts = [] + if custom_mounts: + mounts.extend(custom_mounts) + + env_vars = { + "TRANSFORMERS_OFFLINE": "1", + "TORCH_NCCL_AVOID_RECORD_STREAMS": "1", + "NCCL_NVLS_ENABLE": "0", + "NVTE_DP_AMAX_REDUCE_INTERVAL": "0", + "NVTE_ASYNC_AMAX_REDUCTION": "1", + "NVTE_FUSED_ATTN": "1", + "TOKENIZERS_PARALLELISM": "false", + } + if custom_env_vars: + env_vars |= custom_env_vars + + executor = run.SlurmExecutor( + account=account, + partition=partition, + tunnel=run.SSHTunnel( + user=user, + host=host, + job_dir=remote_job_dir, + ), + exclusive=True, + gres="gpu:8", + nodes=nodes, + ntasks_per_node=devices, + mem="0", + packager=run.GitArchivePackager(subpath="small_language_model_pretraining/nemo", ref="HEAD"), + dependencies=dependencies, + ) + + executor.launcher = None + executor.container_image = container_image + executor.container_mounts = mounts + executor.env_vars = env_vars + executor.retries = retries + executor.time = time + + return executor + +def get_pretrain( + size: str, + nnodes: int, + ngpus_per_node: int, + max_steps: int, + warmup_steps: int, + data_module: run.Config, + max_lr: float = 1e-4, + eval_every: Optional[int] = None, + start_eval_at: Optional[int] = None, + eval_batches: Optional[int] = None, +) -> run.Partial: + + exp_name = size + + pretrain = llm.llama3_8b.pretrain_recipe( + dir="/outputs", + name=exp_name, + num_nodes=nnodes, + num_gpus_per_node=ngpus_per_node + ) + + llama31_config = run.Config(llm.gpt.model.llama.Llama31Config8B) + llama31_config.seq_length = 8192 + pretrain.model.config = llama31_config + + pretrain.trainer.strategy.tensor_model_parallel_size = 1 + pretrain.trainer.strategy.pipeline_model_parallel_size = 1 + pretrain.trainer.strategy.virtual_pipeline_model_parallel_size = 1 # set it back to 7? + pretrain.trainer.strategy.context_parallel_size = 1 + + # Code tracing shows that this is AdamW + pretrain.optim = distributed_fused_adam_with_cosine_annealing( + max_lr=max_lr, + warmup_steps=warmup_steps, + min_lr=max_lr * 0.1 + ) + + precision = run.Config( + nl.MegatronMixedPrecision, + precision="bf16-mixed", + params_dtype=torch.bfloat16, + pipeline_dtype=torch.bfloat16, + autocast_enabled=True, + grad_reduce_in_fp32=False, + fp8="hybrid", + fp8_amax_history_len=4, + fp8_amax_compute_algo='most_recent', + fp8_params=True, + fp8_dot_product_attention=False, + ) + + pretrain.trainer.plugins = precision + + # sets up everything else + pretrain.trainer.max_steps = 1200000 # Hardcoded to fix max_steps for this benchmark + + pretrain.data = data_module + pretrain.trainer.val_check_interval = eval_every / int (os.getenv ("GBS")) + pretrain.trainer.limit_val_batches = eval_batches + pretrain.trainer.limit_test_batches = eval_batches + + pretrain.log.tensorboard = None + pretrain.log.ckpt.every_n_train_steps = None + pretrain.log.ckpt.save_top_k = -1 + pretrain.log.ckpt.save_last = False + pretrain.log.ckpt.always_save_context = False + pretrain.log.ckpt.save_weights_only = False + pretrain.log.ckpt.save_optim_on_train_end = False + pretrain.log.ckpt.save_on_train_epoch_end = False + pretrain.log.ckpt.monitor = "consumed_samples" + pretrain.log.ckpt.mode = "max" + + return exp_name, pretrain + +def get_data( + gbs: int = 288, + mbs: int = 4, + seq_length: Optional[int] = 8192, + tokenizer_path: Optional[str] = "", + seed: Optional[int] = 1234, + use_full_dataset: Optional[bool] = False, +) -> run.Config: + tokenizer = run.Config(AutoTokenizer, pretrained_model_name=tokenizer_path) + + print (f'tokenizer: {tokenizer}') + print (f'use_full_dataset: {use_full_dataset}') + + train_datasets = None + + dataset_path = dataset_path = os.getenv("PREPROCESSED_PATH") + + if use_full_dataset: + train_datasets = sum([["12.5", f"{dataset_path}/c4-train.en_{idx}_text_document"] for idx in range(8)], []) + else: + train_datasets = sum([["10", f"{dataset_path}/c4-train.en_{idx}_text_document"] for idx in [6]], []) + + data_paths = { + "train": train_datasets, + "validation": [ + f"{dataset_path}/c4-validation-91205-samples.en_text_document" + ], + "test": [ + f"{dataset_path}/c4-validation-91205-samples.en_text_document" + ], + } + + return run.Config( + llm.PreTrainingDataModule, + tokenizer=tokenizer, + paths=data_paths, + num_workers=128, # TODO: make it configurable + seq_length=seq_length, + global_batch_size=gbs, + micro_batch_size=mbs, + index_mapping_dir="/npy_index", + seed=seed, + + # Option to reset the position IDs in the dataset at an interval. + reset_position_ids=False, + # Option to reset the attention mask from the dataset. + reset_attention_mask=False, + # Option to enable the EOD mask loss. + eod_mask_loss=False, + # Rampup batch size, should be in format of [start_global_batch_size, batch_size_increment, ramup_samples]. + rampup_batch_size=None, + ) + +def get_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description="Llama3.1 Pretraining") + parser.add_argument("--tag", type=str, help="Optional experiment tag", required=False, default="") + + # Slurm and executor related + slurm_group = parser.add_argument_group("Slurm executor arguments") + slurm_group.add_argument('--user', type=str, required=True, help="Remote cluster SSH user name") + slurm_group.add_argument("--host", type=str, required=True, help="Remote cluster host address") + slurm_group.add_argument("--job_dir", type=str, required=True, help="Remote job directory") + + slurm_group.add_argument("--account", type=str, required=True, help="Account to be used for Slurm job submission") + slurm_group.add_argument("--partition", type=str, required=True, help="Partition to be used for Slurm job submission") + slurm_group.add_argument("--nodes", type=int, required=True, help="Number of nodes to be used") + slurm_group.add_argument("--gpus_per_node", type=int, required=True, help="Number of GPUs per node") + slurm_group.add_argument("--time", type=str, required=True, help="Time limit for the job") + slurm_group.add_argument("--dependencies", nargs="*", help="list of dependencies for the job, dependency type as 'afterok'") # not useful for now + slurm_group.add_argument("--max_retries", type=int, default=0) + slurm_group.add_argument("--run_slurm", action="store_true", help="run in slurm executor instead of locally") + + slurm_group.add_argument( + "--mounts", + type=str, + required=True, + help=( + "Custom mount paths, formatted as a string of :[,:], " + + "and should contain " + + "one path for /output, " + + "NeMo mounted on /opt/NeMo, " + + "dataset path: /workspace/llm/tokenizer.model, /preproc_data, /npy_index" + )) + slurm_group.add_argument("--envvars", type=str, help="Environment variables to be added", default=None) + slurm_group.add_argument("--image", type=str, required=True, help="Container image path, either remote or local") + + model_group = parser.add_argument_group("Model arguments") + model_group.add_argument( + "--size", + type=str, + default="8b", + help="Choose the model to be trained", + choices=[ + "8b", # Llama 3 8B config + ]) + + model_group.add_argument("--initial_ckpt_path", type=str, default=None) + model_group.add_argument("--use_ckpt", action="store_true", help="If set, then resume from the initial checkpoint path") + model_group.add_argument("--resume_from_hf", action="store_true", help="Setting this knob indicates that we are resuming from a weight-only checkpoint") + model_group.add_argument("--ckpt_start_step", type=int, default=0, help="Sets this value to how many steps the resumed checkpoint is already trained on") + model_group.add_argument("--continual_ckpt_path", type=str, default=None, help="Sets this to the path that saves the checkpoint") + model_group.add_argument("--save_ckpt", action="store_true", help="If set, then we save the checkpoint at the end of the experiment") + model_group.add_argument("--tensor_parallel_size", type=int, default=None, help="Set tensor parallelism to the model") + + data_group = parser.add_argument_group("Dataset arguments") + + data_group.add_argument("--gbs", type=int, default=1152, help="Global batch size, should be divisible by PP") + data_group.add_argument("--mbs", type=int, default=1, help="Micro batch size") + data_group.add_argument("--max_lr", type=float, default=1e-4, help="Peak learning rate. Min LR will be 0.1 of max_lr") + data_group.add_argument("--eval_every", type=int, default=46080, help="Evaluate at least every N training sequences") + data_group.add_argument("--start_eval_at", type=int, default=None, help="Start evaluation at N training sequences") + data_group.add_argument("--eval_tokens", type=int, default=1024, help="Evaluate using at least N evaluation sequences") + data_group.add_argument('--max_steps', type=int, default=None, help="Maximum number of steps that each experiment partition will train on. None means no restriction on max steps. ") + data_group.add_argument('--warmup_steps', type=int, default=None, help="Number of steps for LR warmup") + data_group.add_argument("--use_full_dataset", action="store_true", help="If set, then we use the full dataset, instead of the last 256/1024 shards") + data_group.add_argument("--tokenizer_path", type=str, help="Tokenizer path that's used to tokenize the dataset") + + experiment_group = parser.add_argument_group("Experiment management arguments") + experiment_group.add_argument("--dryrun", action="store_true", help="Whether we are launching dryrun or actual runs") + experiment_group.add_argument("--seeds", type=int, nargs="*", default=[], help="random seeds") + experiment_group.add_argument("--num_exps", type=int, default=1) + experiment_group.add_argument("--num_pars", type=int, default=1) + experiment_group.add_argument("--target_log_ppl", type=float, default=5.6) + experiment_group.add_argument("--step_time_atol", type=int, default=1600, help="train step time atol") + + return parser + + +if __name__ == "__main__": + args = get_parser().parse_args() + if args.tag and not args.tag.startswith("-"): + args.tag = "-" + args.tag + + assert not (args.num_pars == 1 and args.continual_ckpt_path is None), "NPar > 1 but a shared checkpoint path is not found" + assert not (not args.save_ckpt and args.num_pars > 1), "multiple experiments are specified but checkpoint is not saved" + + if args.run_slurm: + executor = slurm_executor( + user=args.user, + host=args.host, + remote_job_dir=args.job_dir, + account=args.account, + partition=args.partition, + nodes=args.nodes, + devices=args.gpus_per_node, + time=args.time, + custom_mounts=list(args.mounts.split(",")), + custom_env_vars=({envvar.split("=")[0]: envvar.split("=")[1] for envvar in args.envvars.split(",")} if args.envvars is not None else None), + container_image=args.image, + dependencies=args.dependencies, + retries=args.max_retries, + ) + else: + executor = local_executor( + custom_env_vars=({envvar.split("=")[0]: envvar.split("=")[1] for envvar in args.envvars.split(",")} if args.envvars is not None else None), + devices=args.gpus_per_node, + retries=args.max_retries, + ) + + seq_length = 8192 + + data = get_data( + gbs=args.gbs, + mbs=args.mbs, + seq_length=seq_length, + tokenizer_path=args.tokenizer_path, + seed=1234, # overwritten in each experiments + use_full_dataset=args.use_full_dataset, + ) + + eval_every_n_batches = math.ceil(args.eval_every / (args.gbs)) + eval_batches = math.ceil(args.eval_tokens / (args.gbs)) + if args.start_eval_at is not None: + start_eval_at = math.ceil(args.start_eval_at / args.gbs) + else: + start_eval_at = eval_every_n_batches + + exp_prefix, pretrain = get_pretrain( + max_lr=args.max_lr, + size=args.size, + nnodes=args.nodes, + ngpus_per_node=args.gpus_per_node, + max_steps=args.max_steps, + warmup_steps=args.warmup_steps, + data_module=data, + eval_every=eval_every_n_batches, + start_eval_at=start_eval_at, + eval_batches=eval_batches, + ) + + + # assert args.gbs % pretrain.trainer.strategy.pipeline_model_parallel_size == 0, f"GBS({args.gbs}) should be divisible by PP({pretrain.trainer.strategy.pipeline_model_parallel_size})" + + # Collect all HP configs + from mlperf_logging.mllog import constants + tp = args.tensor_parallel_size or pretrain.trainer.strategy.tensor_model_parallel_size + pp = pretrain.trainer.strategy.pipeline_model_parallel_size + cp = pretrain.trainer.strategy.context_parallel_size + dp = (pretrain.trainer.num_nodes * pretrain.trainer.devices) // (tp * pp * cp) + mini_batch_size = (args.gbs // dp) + grad_accumulation_steps = mini_batch_size // args.mbs + print(f"Parallel settings: {tp=} {pp=} {cp=} {dp=} {mini_batch_size=} {grad_accumulation_steps=}") + # assert False, f"Parallel settings: {tp=} {pp=} {cp=} {pretrain.trainer.num_nodes=} {pretrain.trainer.devices=} {dp=} {args.gbs=} {mini_batch_size=} {grad_accumulation_steps=}" + + configs = { + # HPs + constants.GLOBAL_BATCH_SIZE: args.gbs, + constants.GRADIENT_ACCUMULATION_STEPS: grad_accumulation_steps, + constants.MAX_SEQUENCE_LENGTH: 8192, + constants.EVAL_SAMPLES: args.eval_tokens, + + # Optimizers + constants.OPT_NAME: "adamw", + constants.OPT_BASE_LR: pretrain.optim.config.lr, + constants.OPT_ADAMW_BETA_1: pretrain.optim.config.adam_beta1, + constants.OPT_ADAMW_BETA_2: pretrain.optim.config.adam_beta2, + constants.OPT_ADAMW_EPSILON: pretrain.optim.config.adam_eps, + constants.OPT_ADAMW_WEIGHT_DECAY: pretrain.optim.config.weight_decay, + constants.OPT_GRADIENT_CLIP_NORM: pretrain.optim.config.clip_grad, + + # Schedulers + constants.OPT_END_LR: pretrain.optim.lr_scheduler.min_lr, + constants.OPT_LR_WARMUP_STEPS: pretrain.optim.lr_scheduler.warmup_steps, + constants.OPT_LR_DECAY_STEPS: pretrain.trainer.max_steps - pretrain.optim.lr_scheduler.warmup_steps, + constants.OPT_LR_DECAY_SCHEDULE: "cosine with linear warmup", + } + + # Override config for MLPerf + pretrain.trainer.num_sanity_val_steps = 0 + + run_plugins = [ + plugins.PerfEnvPlugin(), + ] + + exp_prefix = f"{exp_prefix}{args.tag}" + + # Pretrain data index builder + # max steps + pretrain.data.num_train_samples = pretrain.trainer.max_steps * pretrain.data.global_batch_size + print (f'{pretrain.trainer.max_steps=}\n{pretrain.data.global_batch_size=}\n{pretrain.data.num_train_samples=}') + datamodule = pretrain.data.clone() + datamodule.num_dataset_builder_threads = 64 + build_data_index = run.Partial( + build_pretraining_datamodule, + datamodule=datamodule, + trainer_max_steps=pretrain.trainer.max_steps, + trainer_val_check_interval=pretrain.trainer.val_check_interval, + trainer_limit_val_batches=pretrain.trainer.limit_val_batches, + trainer_limit_test_batches=pretrain.trainer.limit_test_batches, + ) + data_index_executor = executor.clone() + data_index_executor.launcher = 'torchrun' + data_index_executor.nodes = 1 + data_index_executor.ntasks_per_node = 1 + data_index_executor.retries = 1 + + static_read_from_path = args.initial_ckpt_path if args.use_ckpt else None + static_write_to_path = args.continual_ckpt_path + static_max_steps = args.max_steps if args.max_steps is not None else None + # Enable this to make static_max_steps not None to enable PreemptiveStop overwrite PL in Callback + static_max_steps = pretrain.trainer.max_steps if static_max_steps is None else static_max_steps + + print (f'{static_max_steps=}') + if not args.save_ckpt: + print (f'Not saving checkpoints') + pretrain.trainer.enable_checkpointing = False + else: + print (f'Saving checkpoints') + + original_callbacks = pretrain.trainer.callbacks + + random_seeds = args.seeds + if len(random_seeds) < args.num_exps: + import random + random_seeds = random_seeds + [random.randint(0, 32767) for _ in range(args.num_exps - len(random_seeds))] + print(f"Missing {args.num_exps - len(random_seeds)} seeds, padding the random seeds to {random_seeds}") + + random_seeds = random_seeds[:args.num_exps] + + for index, seed in enumerate(random_seeds): + # sets the seeds + pretrain.data.seed = seed + build_data_index.datamodule.seed = seed + configs[constants.SEED] = seed + + exp_name = f"{exp_prefix}_{index}_seed_{seed}" + experiment_read_from_path = static_read_from_path + experiment_write_to_path = static_write_to_path + experiment_max_steps = args.ckpt_start_step + + with run.Experiment(exp_name) as exp: + exp.add(build_data_index, executor=data_index_executor, name=f"build_data_index") + + for j in range(args.num_pars): + ending_steps = "" + starting_steps = f"{experiment_max_steps}" + if static_max_steps is not None: + ending_steps = f"-{experiment_max_steps + static_max_steps}-steps" + + print (f'experiment_max_steps: {experiment_max_steps}') + checkpoint_name = "checkpoint" + f"-seed-{seed}-par-{j}{ending_steps}" + experiment_write_to_path = static_write_to_path + "/" + checkpoint_name + + if not (args.resume_from_hf and j == 0): + pretrain.resume = run.Config( + nl.AutoResume, + resume_if_exists=True, + resume_ignore_no_checkpoint=True, + resume_from_path = experiment_read_from_path, + resume_from_directory = experiment_read_from_path, + ) + else: + pretrain.resume = run.Config(nl.AutoResume, restore_config = run.Config(nl.RestoreConfig, path=experiment_read_from_path)) + pretrain.log.ckpt.train_time_interval = None + + if args.save_ckpt: + pretrain.log.ckpt.dirpath = experiment_write_to_path + pretrain.log.ckpt.filename = "checkpoint" + + if static_max_steps is not None: + start_step = experiment_max_steps + experiment_max_steps += static_max_steps + print (f'static_max_steps: {static_max_steps}') + print (f'stop_on_step=experiment_max_steps={experiment_max_steps}') + configs[constants.INIT_CHECKPOINT_STEP] = start_step + pretrain.trainer.callbacks = ( + original_callbacks + [ + run.Config(PreemptiveStop, stop_on_step=experiment_max_steps), + run.Config( + MLPerfCallback, + global_batch_size=args.gbs, + micro_batch_size=args.mbs, + sequence_length=8192, + eval_every=eval_every_n_batches, + init_global_step=start_step, + configs=configs, + ), + ] + ) + + if args.save_ckpt: + pretrain.log.ckpt.every_n_train_steps = experiment_max_steps + pretrain.log.ckpt.save_on_train_epoch_end = False + + try: + print ("control C to skip") + login_info = wandb.login() + print("WandB is logged in.") + pretrain.log.extra_loggers = [ + run.Config( + WandbLogger, + project='llama3.1_8b_training', + name=f'{checkpoint_name}-gbs={args.gbs}-lr={pretrain.optim.config.lr}', + + ), + ] + except: + print("WandB is NOT logged in.") + pretrain.log.extra_loggers = [ + run.Config( + MetricsLogger, + init_global_step=start_step, + global_batch_size=args.gbs, + seq_length=8192, + target_log_ppl=args.target_log_ppl, + train_step_time_atol=args.step_time_atol, + ), + ] + if args.save_ckpt: + pretrain.log.ckpt.every_n_train_steps = experiment_max_steps + pretrain.log.ckpt.save_on_train_epoch_end = False + experiment_read_from_path = experiment_write_to_path + "/checkpoint" + + exp.add( + pretrain, executor=executor, + name=f"{exp_name}_{j}_{starting_steps}{ending_steps}", + plugins=run_plugins + ) + + if args.dryrun: + exp.dryrun() + else: + exp.run(sequential=True, detach=True) diff --git a/small_llm_pretraining/nemo/requirements.txt b/small_llm_pretraining/nemo/requirements.txt new file mode 100644 index 000000000..3dd355195 --- /dev/null +++ b/small_llm_pretraining/nemo/requirements.txt @@ -0,0 +1,9 @@ +git+https://github.com/mlcommons/logging.git@5.0.0-rc2 +git+https://github.com/NVIDIA/mlperf-common.git@68cf1d0d5e3de3351e66abb696d0e2d011aabf47 +huggingface_hub==0.24.0 +transformers==4.43.2 +numpy==1.26.4 +plotly==6.0.0 +nbformat==5.10.4 +kaleido==0.2.1 +redis==5.2.1 \ No newline at end of file diff --git a/small_llm_pretraining/nemo/utils/consolidate_data.sh b/small_llm_pretraining/nemo/utils/consolidate_data.sh new file mode 100644 index 000000000..2b5097cf8 --- /dev/null +++ b/small_llm_pretraining/nemo/utils/consolidate_data.sh @@ -0,0 +1,50 @@ +set -e + +: "${C4_PATH:?C4_PATH not set}" +: "${MERGED_C4_PATH:?MERGED_C4_PATH not set}" +: "${N_VALIDATION_SAMPLES:=91205}" +# defaults the N_VALIDATION_SAMPLES to 91205 +# C4 validation dataset: each sample on average tokenizes to 518 tokens +# thus, to reach 47,185,920 validation tokens, we need to use at least 91205 samples, +# which, after tokenization, will yield 47,186,855 tokens. + +# create softlinks to store each shard before merging +mkdir -p softlinks +# for shard in {0..7}; do +# start=$((shard * 128)) +# end=$((shard * 128 + 127)) +# mkdir -p softlinks/en_$shard +# for ind in $(seq -f "%05g" $start $end); do +# ln -s ${C4_PATH}/c4-train.${ind}-of-01024.json.gz softlinks/en_${shard}/c4-train.${ind}-of-01024.json.gz +# done +# done +for shard in {0..7}; do + start=$((shard * 128)) + end=$((shard * 128 + 127)) + mkdir -p softlinks/en_$shard + for ind in $(seq -f "%05g" $start $end); do + src=${C4_PATH}/c4-train.${ind}-of-01024.json.gz + if [ -f "$src" ]; then + ln -s "$src" softlinks/en_${shard}/ + else + echo "Warning: missing file $src — skipping" >&2 + fi + done +done + +mkdir -p softlinks/en_validation +start=0 +end=7 +for ind in $(seq -f "%05g" $start $end); do + ln -s ${C4_PATH}/c4-validation.${ind}-of-00008.json.gz softlinks/en_validation/c4-validation.${ind}-of-00008.json.gz +done + +# merge +for shard in {0..7}; do + cat softlinks/en_${shard}/*gz > ${MERGED_C4_PATH}/c4-train.en_${shard}.json.gz +done + +cat softlinks/en_validation/*gz > ${MERGED_C4_PATH}/c4-validation.en.json.gz + +# select the first N_VALIDATION_SAMPLES number of samples +zcat ${MERGED_C4_PATH}/c4-validation.en.json.gz | head -n $N_VALIDATION_SAMPLES | gzip > ${MERGED_C4_PATH}/c4-validation-${N_VALIDATION_SAMPLES}-samples.en.json.gz \ No newline at end of file diff --git a/small_llm_pretraining/nemo/utils/download_hf_llama3.sh b/small_llm_pretraining/nemo/utils/download_hf_llama3.sh new file mode 100644 index 000000000..159a50280 --- /dev/null +++ b/small_llm_pretraining/nemo/utils/download_hf_llama3.sh @@ -0,0 +1,2 @@ +huggingface-cli login +huggingface-cli download meta-llama/Llama-3.1-8B --local-dir /data/llama31_8b_ref/model/Llama-3.1-8B-new \ No newline at end of file diff --git a/small_llm_pretraining/nemo/utils/launch_nemo_convert.sh b/small_llm_pretraining/nemo/utils/launch_nemo_convert.sh new file mode 100644 index 000000000..27fdb1d3a --- /dev/null +++ b/small_llm_pretraining/nemo/utils/launch_nemo_convert.sh @@ -0,0 +1,23 @@ +#!/bin/bash +#SBATCH -N 1 +#SBATCH --gpus-per-node 1 +#SBATCH -t 02:00:00 +#SBATCH --mem=0 + +set -e + +: "${CONT_IMAGE_URL:?CONT_IMAGE_URL not set}" +: "${SRC_PATH:?SRC_PATH not set}" +: "${DST_PATH:?DST_PATH not set}" + +working_dir=$(dirname -- ${BASH_SOURCE[0]}) + +if [ ! -d $DST_PATH ]; then + mkdir -p $DST_PATH +fi + +container_maps="${SRC_PATH}:/source,${DST_PATH}:/destination,${working_dir}:/workspace/utils" + +srun --nodes=1 --ntasks-per-node=1 \ +--container-image=$CONT_IMAGE_URL --container-mounts $container_maps --no-container-entrypoint \ +python3 /workspace/utils/convert.py --source /source --destination /destination diff --git a/small_llm_pretraining/nemo/utils/nemo_convert.py b/small_llm_pretraining/nemo/utils/nemo_convert.py new file mode 100644 index 000000000..78d67327d --- /dev/null +++ b/small_llm_pretraining/nemo/utils/nemo_convert.py @@ -0,0 +1,10 @@ +if __name__ == "__main__": + import argparse + from nemo.collections.llm.gpt.model.llama import HFLlamaImporter + parser = argparse.ArgumentParser() + parser.add_argument("--source", default="/source", type=str) + parser.add_argument("--destination", default="/destination", type=str) + args = parser.parse_args() + + importer = HFLlamaImporter(args.source) + importer.apply(args.destination) diff --git a/small_llm_pretraining/nemo/utils/parallel_compress_json_to_gz.sh b/small_llm_pretraining/nemo/utils/parallel_compress_json_to_gz.sh new file mode 100755 index 000000000..2c506708f --- /dev/null +++ b/small_llm_pretraining/nemo/utils/parallel_compress_json_to_gz.sh @@ -0,0 +1,16 @@ +#!/bin/bash +set -e + +: "${C4_PATH:?C4_PATH not set}" + +echo "Starting parallel compression in $C4_PATH..." + +# Use 50% of available CPU cores (adjust -j as needed) +find "$C4_PATH" -maxdepth 1 -name '*.json' | \ + parallel -j$(nproc) ' + echo "Compressing {}" + gzip "{}" +' + +echo "Parallel compression complete!" + diff --git a/small_llm_pretraining/nemo/utils/preprocess.sh b/small_llm_pretraining/nemo/utils/preprocess.sh new file mode 100644 index 000000000..479321a88 --- /dev/null +++ b/small_llm_pretraining/nemo/utils/preprocess.sh @@ -0,0 +1,49 @@ +#!/bin/bash +#SBATCH -N 9 +#SBATCH --gpus-per-node 1 +#SBATCH -t 04:00:00 +#SBATCH --mem=0 + +set -e + +: "${CONT_IMAGE_URL:?CONT_IMAGE_URL not set}" +: "${TOKENIZER_PATH:?TOKENIZER_PATH not set}" +: "${MERGED_C4_PATH:?MERGED_C4_PATH not set}" +: "${PREPROCESSED_PATH:?PREPROCESSED_PATH not set}" + +container_maps="${TOKENIZER_PATH}:/tokenizer,${MERGED_C4_PATH}:/dataset,${PREPROCESSED_PATH}:/outputs" + +# for index in {0..7}; do +# srun --nodes=1 --ntasks-per-node=1 \ +# --container-image=$CONT_IMAGE_URL --container-mounts $container_maps --no-container-entrypoint \ +# python3 /opt/NeMo/scripts/nlp_language_modeling/preprocess_data_for_megatron.py \ +# --input "/dataset/c4-train.en_${index}.json.gz" \ +# --output-prefix "/outputs/c4-train.en_${index}" \ +# --tokenizer-library huggingface --tokenizer-type /tokenizer \ +# --dataset-impl mmap --workers 128 & +# done + +# srun --nodes=1 --ntasks-per-node=1 \ +# --container-image=$CONT_IMAGE_URL --container-mounts $container_maps --no-container-entrypoint \ +# --output preprocess_outputs/dataset_preprocess_validation.out \ +# python3 /opt/NeMo/scripts/nlp_language_modeling/preprocess_data_for_megatron.py \ +# --input "/dataset/c4-validation-91205-samples.en.json.gz" \ +# --output-prefix "/outputs/c4-validation-91205-samples.en" \ +# --tokenizer-library huggingface --tokenizer-type /tokenizer \ +# --dataset-impl mmap --workers 128 & +# wait + +for index in {0..7}; do + python3 /workspace/deps/nemo/scripts/nlp_language_modeling/preprocess_data_for_megatron.py \ + --input "${MERGED_C4_PATH}/c4-train.en_${index}.json.gz" \ + --output-prefix "${PREPROCESSED_PATH}/c4-train.en_${index}" \ + --tokenizer-library huggingface --tokenizer-type ${TOKENIZER_PATH} \ + --dataset-impl mmap --workers 128 & +done + + python3 /workspace/deps/nemo/scripts/nlp_language_modeling/preprocess_data_for_megatron.py \ + --input "${MERGED_C4_PATH}/c4-validation-91205-samples.en.json.gz" \ + --output-prefix "${PREPROCESSED_PATH}/c4-validation-91205-samples.en" \ + --tokenizer-library huggingface --tokenizer-type ${TOKENIZER_PATH} \ + --dataset-impl mmap --workers 128 & +wait