diff --git a/examples/qwen_moe_nemo/train_without_nemorun/train.py b/examples/qwen_moe_nemo/train_without_nemorun/train.py new file mode 100644 index 0000000..7b23baf --- /dev/null +++ b/examples/qwen_moe_nemo/train_without_nemorun/train.py @@ -0,0 +1,215 @@ +import torch +from nemo import lightning as nl +import nemo_run as run +from nemo.collections import llm +from megatron.core.optimizer import OptimizerConfig +import pytorch_lightning as pl +from typing import List, Optional +from torch.utils.data import DataLoader, Dataset +from datasets import load_dataset +from transformers import AutoTokenizer + +from nemo.utils.import_utils import safe_import +from nemo.collections.nlp.modules.common.tokenizer_utils import get_nmt_tokenizer + +_, HAVE_TE = safe_import("transformer_engine") + + +class CustomDataModule(pl.LightningDataModule): + def __init__( + self, + dataset_name: str = "tatsu-lab/alpaca", + seq_length: int = 2048, + tokenizer: Optional["TokenizerSpec"] = None, # autotokenizer? + micro_batch_size: int = 4, + global_batch_size: int = 8, + rampup_batch_size: Optional[List[int]] = None, + num_train_samples: int = 10_000_000, + num_val_samples: int = 10_000, + num_test_samples: int = 10_000, + num_workers: int = 8, + pin_memory: bool = True, + persistent_workers: bool = False, + create_attention_mask: bool = False, + vocab_file: Optional[str] = None, + merges_file: Optional[str] = None, + ): + super().__init__() + self.dataset_name = dataset_name + self.seq_length = seq_length + self.micro_batch_size = micro_batch_size + self.global_batch_size = global_batch_size + self.num_train_samples = num_train_samples + self.num_val_samples = num_val_samples + self.num_test_samples = num_test_samples + self.num_workers = num_workers + self.pin_memory = pin_memory + self.persistent_workers = persistent_workers + self.create_attention_mask = create_attention_mask + # self.create_attention_mask = create_attention_mask or not HAVE_TE + + if tokenizer is None: + from nemo.collections.nlp.modules.common.tokenizer_utils import get_nmt_tokenizer + self.tokenizer = get_nmt_tokenizer( + "megatron", "GPT2BPETokenizer", vocab_file=vocab_file, merges_file=merges_file + ) + else: + self.tokenizer = tokenizer + + # self.data_sampler = MegatronDataSampler( + # seq_len=self.seq_length, + # micro_batch_size=self.micro_batch_size, + # global_batch_size=self.global_batch_size, + # rampup_batch_size=rampup_batch_size, + # ) + + def setup(self, stage: str = "") -> None: + """ + Setup the data module. + """ + self._train_ds = load_dataset(self.dataset_name, split="train") + self._validation_ds = load_dataset(self.dataset_name, split="validation") + self._test_ds = load_dataset(self.dataset_name, split="test") + + def train_dataloader(self): + """ + Get the train dataloader. + """ + if not hasattr(self, "_train_ds"): + self.setup() + return self._create_dataloader(self._train_ds) + + def val_dataloader(self): + """ + Get the validation dataloader. + """ + if not hasattr(self, "_validation_ds"): + self.setup() + return self._create_dataloader(self._validation_ds) + + def test_dataloader(self): + """ + Get the test dataloader. + """ + if not hasattr(self, "_test_ds"): + self.setup() + return self._create_dataloader(self._test_ds) + + def _create_dataloader(self, dataset, **kwargs) -> DataLoader: + return DataLoader( + dataset, + num_workers=self.num_workers, + pin_memory=self.pin_memory, + persistent_workers=self.persistent_workers, + collate_fn=dataset.collate_fn, + **kwargs, + ) + + def reconfigure_limit_batches(self): + """ + Reconfigure trainer.limit_train_batches and trainer.limit_val_batches in terms of num of microbatches. + """ + from nemo.collections.llm.gpt.data.utils import _reconfigure_limit_batches + + # Override limit_train_batches in terms of num of microbatches + self.trainer.limit_train_batches = _reconfigure_limit_batches(self.trainer.limit_train_batches, self._train_ds) + # Override limit_val_batches to be a multiple of num microbatches to prevent val_step from exiting + # in between a step + self.trainer.limit_val_batches = _reconfigure_limit_batches( + self.trainer.limit_val_batches, self._validation_ds + ) + + try: + from megatron.core.num_microbatches_calculator import get_num_microbatches + + except (ImportError, ModuleNotFoundError): + from apex.transformer.pipeline_parallel.utils import get_num_microbatches + + # Override num sanity steps to be a multiple of num of microbatches + self.trainer.num_sanity_val_steps *= get_num_microbatches() + +if __name__ == "__main__": + seq_length = 2048 + global_batch_size = 16 + + ## setup the dummy dataset + # data = llm.MockDataModule(seq_length=seq_length, global_batch_size=global_batch_size) + tokenizer = get_nmt_tokenizer( + "megatron", "GPT2BPETokenizer" + ) + data = llm.HFDatasetDataModule(path_or_dataset="nvidia/OpenMathInstruct-1", tokenizer=tokenizer) + # data = CustomDataModule( + # dataset_name="tatsu-lab/alpaca", + # seq_length=seq_length, + # global_batch_size=global_batch_size, + # # tokenizer=AutoTokenizer.from_pretrained("Qwen/Qwen3-30B-A3B"), + # ) + + ## initialize a small GPT model + gpt_config = llm.GPTConfig( + num_layers=6, + hidden_size=384, + ffn_hidden_size=1536, + num_attention_heads=6, + seq_length=seq_length, + init_method_std=0.023, + hidden_dropout=0.1, + attention_dropout=0.1, + layernorm_epsilon=1e-5, + make_vocab_size_divisible_by=128, + ) + + + model = llm.GPTModel(gpt_config, tokenizer=data.tokenizer) + # model = llm.GPTModel(gpt_config, tokenizer=tokenizer) + + ## initialize the strategy + strategy = nl.MegatronStrategy( + tensor_model_parallel_size=1, + pipeline_model_parallel_size=1, + pipeline_dtype=torch.bfloat16, + ) + + ## setup the optimizer + opt_config = OptimizerConfig( + optimizer='adam', + lr=6e-4, + bf16=True, + ) + opt = nl.MegatronOptimizerModule(config=opt_config) + + trainer = nl.Trainer( + devices=2, ## you can change the number of devices to suit your setup + max_steps=200, + accelerator="gpu", + strategy=strategy, + plugins=nl.MegatronMixedPrecision(precision="bf16-mixed"), + ) + + nemo_logger = nl.NeMoLogger( + log_dir="test_logdir", ## logs and checkpoints will be written here + ) + print("Done initializing logger ...") + print("Beginning training...") + + llm.train( + model=model, + data=data, + trainer=trainer, + log=nemo_logger, + tokenizer='data', + optim=opt, + ) + + # recipe = run.Partial( + # llm.train, + # model=model, + # data=data, + # trainer=trainer, + # log=nemo_logger, + # tokenizer='data', + # optim=opt, + # ) + + # run.run(recipe, executor=run.LocalExecutor()) + diff --git a/examples/qwen_moe_nemo/training/config.py b/examples/qwen_moe_nemo/training/config.py new file mode 100644 index 0000000..f3575d8 --- /dev/null +++ b/examples/qwen_moe_nemo/training/config.py @@ -0,0 +1,41 @@ +from truss_train import definitions +from truss.base import truss_config + +BASE_IMAGE = "nvcr.io/nvidia/nemo:25.07" + +training_runtime = definitions.Runtime( + start_commands=[ + "./run.sh" + ], + environment_variables={ + "HF_TOKEN": definitions.SecretReference(name="hf_access_token"), + # "HF_HOME": "/root/.cache/user_artifacts/hf_cache", + # "WANDB_API_KEY": definitions.SecretReference(name="wandb_api_key"), + }, + cache_config=definitions.CacheConfig( + enabled=False, + ), + checkpointing_config=definitions.CheckpointingConfig( + enabled=True, + checkpoint_path="/tmp/training_checkpoints", + ), +) + +training_compute = definitions.Compute( + accelerator=truss_config.AcceleratorSpec( + accelerator=truss_config.Accelerator.H100, + count=8, + ), + node_count=1, +) + +training_job = definitions.TrainingJob( + image=definitions.Image(base_image=BASE_IMAGE), + compute=training_compute, + runtime=training_runtime +) + +training_project = definitions.TrainingProject( + name="Nemo template", + job=training_job +) \ No newline at end of file diff --git a/examples/qwen_moe_nemo/training/run.sh b/examples/qwen_moe_nemo/training/run.sh new file mode 100644 index 0000000..88675a8 --- /dev/null +++ b/examples/qwen_moe_nemo/training/run.sh @@ -0,0 +1,9 @@ +#!/usr/bin/env bash +set -euo pipefail + +export NEMO_CACHE_DIR=$BT_RW_CACHE_DIR +export NEMO_MODELS_CACHE=$NEMO_CACHE_DIR/nemo_models + + +nemo llm import model=qwen3_30b_a3b source="hf://Qwen/Qwen3-30B-A3B" -y +nemo llm finetune --factory "qwen3_30b_a3b(peft_scheme=none)" -y # uses dummy dataset \ No newline at end of file diff --git a/examples/qwen_moe_nemo/training/train.py b/examples/qwen_moe_nemo/training/train.py new file mode 100644 index 0000000..2ed2615 --- /dev/null +++ b/examples/qwen_moe_nemo/training/train.py @@ -0,0 +1,61 @@ +from nemo.collections import llm +import nemo_run as run +from nemo.collections.llm.gpt.data.hf_dataset import HFDatasetDataModule + +# import lightning.pytorch as pl # data module has to be pl.LightningDataModule + +# """ +# export NEMO_CACHE_DIR=$BT_RW_CACHE_DIR +# export NEMO_MODELS_CACHE=$NEMO_CACHE_DIR/nemo_models + +# Examples: +# >>> from nemo.collections import llm +# >>> from nemo import lightning as nl +# >>> model = llm.MistralModel() +# >>> data = llm.SquadDataModule(seq_length=4096, global_batch_size=16, micro_batch_size=2) +# >>> precision = nl.MegatronMixedPrecision(precision="bf16-mixed") +# >>> trainer = nl.Trainer(strategy=nl.MegatronStrategy(tensor_model_parallel_size=2), plugins=precision) +# >>> llm.finetune(model, data, trainer, peft=llm.peft.LoRA()]) +# """ + +model = llm.Qwen3Model(llm.Qwen3Config30B_A3B()) + +llm.import_ckpt(model=model, source="hf://Qwen/Qwen3-30B-A3B") + + + +# recipe = run.Partial( +# llm.finetune, +# model=model, +# trainer=default_finetune_trainer( +# num_nodes=num_nodes, +# num_gpus_per_node=num_gpus_per_node, +# ), +# data=HFDatasetDataModule(path="tatsu-lab/alpaca", global_batch_size=16, micro_batch_size=2), +# log=default_finetune_log(dir=dir, name=name, tensorboard_logger=tensorboard_logger(name=name)), +# optim=distributed_fused_adam_with_cosine_annealing(max_lr=1e-4, min_lr=0, warmup_steps=50, adam_beta2=0.98), +# resume=nemo_resume(resume_path), +# tokenizer=tokenizer, +# ) + + + +# return recipe + +recipe = llm.qwen3_30b_a3b.finetune_recipe( + name="qwen3_30b_a3b_ft_nolora", + dir="checkpoints", + num_nodes=1, + num_gpus_per_node=8, + # peft_scheme='lora', + packed_sequence=False + ) + +# recipe.data = HFDatasetDataModule(path_or_dataset="tatsu-lab/alpaca", global_batch_size=16, micro_batch_size=2) +# recipe.data = HFDatasetDataModule(path_or_dataset="tatsu-lab/alpaca") + +recipe.data = llm.HFDatasetDataModule(path_or_dataset="nvidia/OpenMathInstruct-1") + + +run.run(recipe, executor=run.LocalExecutor()) + diff --git a/examples/qwen_moe_nemo/training_distil_qwen2.5/bespoke.py b/examples/qwen_moe_nemo/training_distil_qwen2.5/bespoke.py new file mode 100644 index 0000000..7598576 --- /dev/null +++ b/examples/qwen_moe_nemo/training_distil_qwen2.5/bespoke.py @@ -0,0 +1,147 @@ +import json +import shutil +from typing import TYPE_CHECKING, Any, Dict, List, Optional + +# import numpy as np +from datasets import load_dataset + +from nemo.collections.llm.gpt.data.core import get_dataset_root +from nemo.collections.llm.gpt.data.fine_tuning import FineTuningDataModule +from nemo.lightning.io.mixin import IOMixin +from nemo.utils import logging + +from functools import lru_cache + +from nemo.collections.llm.gpt.data.core import create_sft_dataset + +if TYPE_CHECKING: + from nemo.collections.common.tokenizers import TokenizerSpec + from nemo.collections.llm.gpt.data.packed_sequence import PackedSequenceSpecs + + +class BespokeDataModule(FineTuningDataModule, IOMixin): + """A data module for fine-tuning on the Bespoke dataset. + + This class inherits from the `FineTuningDataModule` class and is specifically designed for fine-tuning models on the + "bespokelabs/Bespoke-Stratos-17k" dataset. It handles data download, preprocessing, splitting, and preparing the data + in a format suitable for training, validation, and testing. + + Args: + force_redownload (bool, optional): Whether to force re-download the dataset even if it exists locally. Defaults to False. + delete_raw (bool, optional): Whether to delete the raw downloaded dataset after preprocessing. Defaults to True. + See FineTuningDataModule for the other args + """ + + def __init__( + self, + seq_length: int = 2048, + tokenizer: Optional["TokenizerSpec"] = None, + micro_batch_size: int = 4, + global_batch_size: int = 8, + rampup_batch_size: Optional[List[int]] = None, + force_redownload: bool = False, + delete_raw: bool = True, + seed: int = 1234, + memmap_workers: int = 1, + num_workers: int = 8, + pin_memory: bool = True, + persistent_workers: bool = False, + packed_sequence_specs: Optional["PackedSequenceSpecs"] = None, + dataset_kwargs: Optional[Dict[str, Any]] = None, + dataset_root: str = "./bespoke", + ): + self.force_redownload = force_redownload + self.delete_raw = delete_raw + + super().__init__( + dataset_root=dataset_root, + seq_length=seq_length, + tokenizer=tokenizer, + micro_batch_size=micro_batch_size, + global_batch_size=global_batch_size, + rampup_batch_size=rampup_batch_size, + seed=seed, + memmap_workers=memmap_workers, + num_workers=num_workers, + pin_memory=pin_memory, + persistent_workers=persistent_workers, + packed_sequence_specs=packed_sequence_specs, + dataset_kwargs=dataset_kwargs, + ) + + def prepare_data(self) -> None: + # if train file is specified, no need to do anything + if not self.train_path.exists() or self.force_redownload: + dset = self._download_data() + self._preprocess_and_split_data(dset) + super().prepare_data() + + def _download_data(self): + logging.info(f"Downloading {self.__class__.__name__}...") + return load_dataset( + "bespokelabs/Bespoke-Stratos-17k", + cache_dir=str(self.dataset_root), + download_mode="force_redownload" if self.force_redownload else None, + ) + + def _preprocess_and_split_data(self, dset, train_ratio: float = 0.80, val_ratio: float = 0.15): + logging.info(f"Preprocessing {self.__class__.__name__} to jsonl format and splitting...") + test_ratio = 1 - train_ratio - val_ratio + save_splits = {} + dataset = dset.get('train') + split_dataset = dataset.train_test_split(test_size=val_ratio + test_ratio, seed=self.seed) + split_dataset2 = split_dataset['test'].train_test_split( + test_size=test_ratio / (val_ratio + test_ratio), seed=self.seed + ) + save_splits['training'] = split_dataset['train'] + save_splits['validation'] = split_dataset2['train'] + save_splits['test'] = split_dataset2['test'] + + print("len training: ", len(save_splits['training'])) + print("len validation: ", len(save_splits['validation'])) + print("len test: ", len(save_splits['test'])) + + for split_name, dataset in save_splits.items(): + output_file = self.dataset_root / f"{split_name}.jsonl" + with output_file.open("w", encoding="utf-8") as f: + for example in dataset: + + conversations = example["conversations"] + + for conversation in conversations: + if conversation["from"] == "user": + conversation["from"] = "User" + elif conversation["from"] == "assistant": + conversation["from"] = "Assistant" + else: + raise ValueError(f"Unknown role: {conversation['role']}") + + example["mask"] = "User" + example["type"] = "VALUE_TO_TEXT" + + f.write(json.dumps(example) + "\n") + + logging.info(f"{split_name} split saved to {output_file}") + + if self.delete_raw: + for p in self.dataset_root.iterdir(): + if p.is_dir(): + shutil.rmtree(p) + elif '.jsonl' not in str(p.name): + p.unlink() + + @lru_cache + def _create_dataset(self, path, pack_metadata_path=None, is_test=False, **kwargs): + # pylint: disable=C0115,C0116 + return create_sft_dataset( + path, + tokenizer=self.tokenizer, + seq_length=(self.seq_length if is_test or self.packed_sequence_size <= 0 else self.packed_sequence_size), + memmap_workers=self.memmap_workers, + seed=self.seed, + chat=True, + is_test=is_test, + pack_metadata_file_path=None, # packing is not supported + pad_cu_seqlens=False, + **kwargs, + ) diff --git a/examples/qwen_moe_nemo/training_distil_qwen2.5/train.py b/examples/qwen_moe_nemo/training_distil_qwen2.5/train.py new file mode 100644 index 0000000..e1189d8 --- /dev/null +++ b/examples/qwen_moe_nemo/training_distil_qwen2.5/train.py @@ -0,0 +1,143 @@ +import nemo_run as run +import lightning.pytorch as pl +from nemo.collections import llm + +# llm.import_ckpt is the nemo2 API for converting Hugging Face checkpoint to NeMo format +# example python usage: +# llm.import_ckpt(model=llm.llama3_8b.model(), source="hf://meta-llama/Meta-Llama-3-8B") +# +# We use run.Partial to configure this function +def configure_checkpoint_conversion(): + return run.Partial( + llm.import_ckpt, + model=llm.qwen2_7b.model(), + source="hf://Qwen/Qwen2.5-7B-Instruct", + overwrite=True, + ) + +# configure your function +import_ckpt = configure_checkpoint_conversion() +# define your executor +local_executor = run.LocalExecutor() + +# run your experiment +run.run(import_ckpt, executor=local_executor) + +### Dataset +from bespoke import BespokeDataModule + +def bespoke() -> run.Config[pl.LightningDataModule]: + return run.Config(BespokeDataModule, seq_length=16384, micro_batch_size=1, global_batch_size=32, num_workers=1) + +import nemo_run as run +from nemo import lightning as nl +from nemo.collections import llm +from megatron.core.optimizer import OptimizerConfig +import torch +import lightning.pytorch as pl +from pathlib import Path +from nemo.collections.llm.recipes.precision.mixed_precision import bf16_mixed + + +# Configure the trainer +# we use 4 GPUs for training and set the max_steps to 300. +def trainer() -> run.Config[nl.Trainer]: + strategy = run.Config( + nl.MegatronStrategy, + tensor_model_parallel_size=4, + ) + trainer = run.Config( + nl.Trainer, + devices=4, + max_steps=300, + accelerator="gpu", + strategy=strategy, + plugins=bf16_mixed(), + log_every_n_steps=1, + limit_val_batches=0, + val_check_interval=0, + num_sanity_val_steps=0, + ) + return trainer + +# Configure the logger +# Here, we configure the log interval to 100 steps and save the model every 100 steps. you can change these parameters as needed. +def logger() -> run.Config[nl.NeMoLogger]: + ckpt = run.Config( + nl.ModelCheckpoint, + save_last=True, + every_n_train_steps=100, + monitor="reduced_train_loss", + save_top_k=1, + save_on_train_epoch_end=True, + save_optim_on_train_end=True, + ) + + return run.Config( + nl.NeMoLogger, + name="qwen_sft", + log_dir="./results", + use_datetime_version=False, + ckpt=ckpt, + wandb=None + ) + + +# Configure the optimizer +# We use the distributed Adam optimizer and pass in the OptimizerConfig. +def adam_with_cosine_annealing() -> run.Config[nl.OptimizerModule]: + opt_cfg = run.Config( + OptimizerConfig, + optimizer="adam", + lr=2e-5, + adam_beta2=0.98, + use_distributed_optimizer=True, + clip_grad=1.0, + bf16=True, + ) + return run.Config( + nl.MegatronOptimizerModule, + config=opt_cfg + ) + +# Configure the model +# We use Qwen2Config7B to configure the model. +def qwen() -> run.Config[pl.LightningModule]: + return run.Config(llm.Qwen2Model, config=run.Config(llm.Qwen2Config7B)) + +# Configure the resume +def resume() -> run.Config[nl.AutoResume]: + return run.Config( + nl.AutoResume, + restore_config=run.Config(nl.RestoreConfig, + path="nemo://Qwen/Qwen2.5-7B-Instruct" + ), + resume_if_exists=True, + ) + +def configure_finetuning_recipe(): + return run.Partial( + llm.finetune, + model=qwen(), + trainer=trainer(), + data=bespoke(), + log=logger(), + optim=adam_with_cosine_annealing(), + resume=resume(), + ) + +def local_executor_torchrun(nodes: int = 1, devices: int = 4) -> run.LocalExecutor: + # Env vars for jobs are configured here + env_vars = { + "TORCH_NCCL_AVOID_RECORD_STREAMS": "1", + "NCCL_NVLS_ENABLE": "0", + "NVTE_DP_AMAX_REDUCE_INTERVAL": "0", + "NVTE_ASYNC_AMAX_REDUCTION": "1", + } + + executor = run.LocalExecutor(ntasks_per_node=devices, launcher="torchrun", env_vars=env_vars) + + return executor + +if __name__ == '__main__': + run.run(configure_finetuning_recipe(), executor=local_executor_torchrun())