run_train.py

"""
Nanotron training script.

Usage:
```
export CUDA_DEVICE_MAX_CONNECTIONS=1 # important for some distributed operations
torchrun --nproc_per_node=8 run_train.py --config-file examples/config_tiny_llama.yaml
```
"""
import argparse
import time
from pprint import pformat
from typing import Dict, Optional, cast

import nanotron.distributed as dist
from nanotron import logging
from nanotron.config import (
    DataArgs,
    DatasetStageArgs,
    NanosetDatasetsArgs,
    PretrainDatasetsArgs,
    Qwen2Config,
    SFTDatasetsArgs,
)
from nanotron.data.dataloader import (
    dummy_infinite_data_generator,
    get_train_dataloader,
)
from nanotron.data.processing import (
    clm_process,
    get_datasets,
)
from nanotron.data.sft_processing import prepare_sft_dataset
from nanotron.helpers import (
    compute_remain_train_steps_of_a_data_stage_from_ckp,
    get_consumed_train_samples_of_a_data_stage_from_ckp,
)
from nanotron.logging import log_rank
from nanotron.parallel.pipeline_parallel.utils import get_input_output_pp_ranks
from nanotron.sanity_checks import sanity_check_dataloader
from nanotron.trainer import DistributedTrainer
from nanotron.utils import main_rank_first
from torch.utils.data import DataLoader
from nanotron.trainer import DataStageMetadata
from collections import defaultdict
try:
    from huggingface_hub import __version__ as hf_hub_version
    from transformers import AutoTokenizer
    from transformers import __version__ as tf_version
except ImportError:
    hf_hub_version = None
    tf_version = None

logger = logging.get_logger(__name__)

# import lovely_tensors as lt

# lt.monkey_patch()


def get_dataloader_from_data_stage(
    trainer: DistributedTrainer,
    data: DataArgs,
    consumed_train_samples_stage: int,
    consumed_tokens_per_dataset_folder: Dict[str, int],
    last_stages_consumed_tokens_per_dataset_folder: Dict[str, int],
    num_remaining_train_steps: int,
    sanity_check_dataloader_interval: Optional[int] = None,
):
    """
    Returns a dataloader for a given data stage.

    data: The data configuration for the current stage.
    consumed_train_samples_stage: The number of samples consumed by the model in the this stage (each stage starts from zero).
    consumed_tokens_per_dataset_folder: The number of tokens consumed by the model in previous stages to avoid reseeing them, because the sampler has restarted for this stage.
    num_remaining_train_steps: The number of remaining training steps for this stage.
    """
    assert consumed_train_samples_stage >= 0, "consumed_train_samples_stage should be greater than 0"
    assert num_remaining_train_steps >= 0, "num_remaining_train_steps should be greater than 0"

    # First, we need to know which ranks to feed the dataloader to
    input_pp_rank, output_pp_rank = get_input_output_pp_ranks(model=trainer.model)

    # Case 1: Dummy data generator
    if data.dataset is None:
        log_rank("Using dummy data generator", logger=logger, level=logging.INFO, rank=0)
        dataloader = dummy_infinite_data_generator(
            micro_batch_size=trainer.micro_batch_size,
            sequence_length=trainer.sequence_length,
            input_pp_rank=input_pp_rank,
            output_pp_rank=output_pp_rank,
            vocab_size=trainer.model_config.vocab_size,
            seed=data.seed,
            parallel_context=trainer.parallel_context,
            use_position_ids=isinstance(
                trainer.model_config, Qwen2Config
            ),  # Simulate packed sequences to test SFT or inference
            cp_pg=trainer.parallel_context.cp_pg,
        )()

    # Case 2: HuggingFace datasets
    elif isinstance(data.dataset, PretrainDatasetsArgs) or isinstance(data.dataset, SFTDatasetsArgs):
        log_rank("Using `datasets` library", logger=logger, level=logging.INFO, rank=0)
        tokenizer_path = trainer.config.tokenizer.tokenizer_name_or_path
        log_rank(
            f"Loading tokenizer from {tokenizer_path} and transformers/hf_hub versions {tf_version, hf_hub_version}",
            logger=logger,
            level=logging.INFO,
            rank=0,
        )

        # We need to the 1st device to process dataset and cache it, then other devices load from cache
        with main_rank_first(trainer.parallel_context.world_pg):
            # TODO @nouamanetazi: this may timeout before 1st device finishes processing dataset. Can we have a ctxmanager to modify timeout?
            # TODO: generalise to include  for validation/test splits

            # We load the raw dataset
            raw_dataset = get_datasets(
                hf_dataset_or_datasets=data.dataset.hf_dataset_or_datasets,
                hf_dataset_config_name=data.dataset.hf_dataset_config_name,
                splits=data.dataset.hf_dataset_splits,
            )["train"]

            tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
            tokenizer.padding_side = "left"
            sequence_sep_tokens = [tokenizer.bos_token, tokenizer.eos_token, tokenizer.pad_token, tokenizer.unk_token]
            # assert bos or eos are present
            assert (
                tokenizer.bos_token is not None or tokenizer.eos_token is not None
            ), f"Tokenizer must have either bos or eos token, but found none for {tokenizer_path}"

            # Check that tokenizer's vocab size is smaller than the model's vocab size
            assert (
                tokenizer.vocab_size <= trainer.model_config.vocab_size
            ), f"Tokenizer's vocab size ({tokenizer.vocab_size}) is larger than the model's vocab size ({trainer.model_config.vocab_size})"

            # Different processing for SFT vs pretraining
            if isinstance(data.dataset, SFTDatasetsArgs):
                # For SFT, use the dedicated prepare_sft_dataset function
                # Get optional debug parameter to limit dataset size (for faster development)
                debug_max_samples = getattr(data.dataset, "debug_max_samples", None)

                # Process the dataset using our dedicated SFT processing module
                train_dataset = prepare_sft_dataset(
                    raw_dataset=raw_dataset,
                    tokenizer=tokenizer,
                    trainer_sequence_length=trainer.sequence_length,
                    debug_max_samples=debug_max_samples,
                    num_proc=data.dataset.dataset_processing_num_proc_per_process,
                )
            else:
                # For pretraining, use existing CLM processing
                train_dataset = clm_process(
                    raw_dataset=raw_dataset,
                    tokenizer=tokenizer,
                    text_column_name=data.dataset.text_column_name,
                    dataset_processing_num_proc_per_process=data.dataset.dataset_processing_num_proc_per_process,
                    dataset_overwrite_cache=data.dataset.dataset_overwrite_cache,
                    sequence_length=trainer.sequence_length,
                )

            # We load the processed dataset on the ranks requiring it
            dataloader = get_train_dataloader(
                train_dataset=train_dataset,
                sequence_length=trainer.sequence_length,
                parallel_context=trainer.parallel_context,
                input_pp_rank=input_pp_rank,
                output_pp_rank=output_pp_rank,
                micro_batch_size=trainer.micro_batch_size,
                consumed_train_samples_stage=consumed_train_samples_stage,
                dataloader_num_workers=data.num_loading_workers,
                seed_worker=data.seed,
                dataloader_drop_last=True,
                use_position_ids=isinstance(trainer.model_config, Qwen2Config),
                sequence_sep_tokens=sequence_sep_tokens,  # Used to generate position ids
            )

            # Check if we have enough samples for train_steps
            total_tokens_dataset = len(dataloader.dataset) * trainer.sequence_length
            num_tokens_needed_for_training = (
                num_remaining_train_steps * trainer.global_batch_size * trainer.sequence_length
            )
            assert num_tokens_needed_for_training <= total_tokens_dataset, (
                f"Dataset is too small for steps ({total_tokens_dataset} < {num_tokens_needed_for_training}), "
                f"Try train_steps<={len(dataloader.dataset) // trainer.global_batch_size + trainer.iteration_step}"
            )

    # Case 3: Nanosets
    elif isinstance(data.dataset, NanosetDatasetsArgs):
        log_rank("Using TokenizedBytes Dataloader", logger=logger, level=logging.INFO, rank=0)
        from nanotron.data.tokenized_bytes import get_tb_dataloader, get_tb_datasets

        tokenizer_path = trainer.config.tokenizer.tokenizer_name_or_path
        tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
        assert (
            len(tokenizer) == trainer.model_config.vocab_size
        ), f"Tokenizer vocab size ({len(tokenizer)}) does not match model config vocab size ({trainer.model_config.vocab_size}). "
        log_rank(
            f"[TokenizedBytes] Creating TokenizedBytes with {len(data.dataset.dataset_folder)} dataset folders and {trainer.config.tokens.train_steps * trainer.global_batch_size} train samples",
            logger=logger,
            level=logging.INFO,
            rank=0,
        )
        if tokenizer.pad_token is None:
            tokenizer.pad_token = tokenizer.eos_token
        tokenizer.padding_side = "left"

        start_time = time.time()
        train_dataset, data_log = get_tb_datasets(
            config=data.dataset,
            global_batch_size=trainer.global_batch_size,
            sequence_length=trainer.sequence_length,
            train_steps=trainer.config.tokens.train_steps,
            current_iteration=trainer.iteration_step,
            parallel_context=trainer.parallel_context,
            shuffle=data.dataset.shuffle_files,
            eos_token_id=tokenizer.eos_token_id,
            seed=data.seed,
            consumed_samples=consumed_train_samples_stage,
            consumed_tokens_per_dataset_folder=consumed_tokens_per_dataset_folder,
            last_stages_consumed_tokens_per_dataset_folder=last_stages_consumed_tokens_per_dataset_folder,
        )
        dataloader = get_tb_dataloader(
            dataset=train_dataset,
            sequence_length=trainer.sequence_length,
            micro_batch_size=trainer.micro_batch_size,
            global_batch_size=trainer.global_batch_size,
            num_workers=data.num_loading_workers,
            cfg=data.dataset,
            consumed_samples=consumed_train_samples_stage,
            num_samples=trainer.config.tokens.train_steps * trainer.global_batch_size, # TODO: this overshoots what's needed by the current stage, but it doesnt matter?
            parallel_context=trainer.parallel_context,
            input_pp_rank=input_pp_rank,
            output_pp_rank=output_pp_rank,
            dataloader_drop_last=True,
            dataloader_pin_memory=True,
            use_position_ids=isinstance(trainer.model_config, Qwen2Config),
            use_doc_masking=getattr(trainer.model_config, "_use_doc_masking", None),
        )
        log_rank(
            f"[TokenizedBytes] Time taken to create TokenizedBytes: {time.strftime('%M:%S', time.gmtime(time.time() - start_time))} (MM:SS)",
            logger=logger,
            level=logging.INFO,
            rank=0,
        )
        dist.barrier()

        # Create Nanoset
        # from nanotron.data.nanoset import Nanoset

        # with main_rank_first(trainer.parallel_context.world_pg):
        #     tokenizer_path = trainer.config.tokenizer.tokenizer_name_or_path
        #     tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
        #     eos_token_id = tokenizer.eos_token_id
        #     assert (
        #         eos_token_id is not None or data.dataset.return_positions is False
        #     ), "Tokenizer must have an eos token if return_positions is True"
        #     log_rank(
        #         f"[Nanoset] Creating Nanoset with {len(data.dataset.dataset_folder)} dataset folders and {trainer.config.tokens.train_steps * trainer.global_batch_size} train samples",
        #         logger=logger,
        #         level=logging.INFO,
        #         rank=0,
        #     )
        #     start_time = time.time()
        #     train_dataset = Nanoset(
        #         dataset_folders=data.dataset.dataset_folder,
        #         sequence_length=trainer.sequence_length,
        #         token_size=data.dataset.token_size_in_bytes,
        #         train_split_num_samples=trainer.config.tokens.train_steps * trainer.global_batch_size,
        #         dataset_weights=data.dataset.dataset_weights,
        #         random_seed=data.seed,
        #         return_positions=data.dataset.return_positions,
        #         eos_token_id=eos_token_id,
        #     )
        #     end_time = time.time()
        #     log_rank(
        #         f"[Nanoset] Time taken to create Nanoset: {time.strftime('%M:%S', time.gmtime(end_time - start_time))} (MM:SS)",
        #         logger=logger,
        #         level=logging.INFO,
        #         rank=0,
        #     )
        # # Prepare dataloader
        # train_dataloader = build_nanoset_dataloader(
        #     train_dataset,
        #     trainer.sequence_length,
        #     parallel_context=trainer.parallel_context,
        #     input_pp_rank=input_pp_rank,
        #     output_pp_rank=output_pp_rank,
        #     micro_batch_size=trainer.micro_batch_size,
        #     consumed_train_samples=consumed_train_samples,
        #     dataloader_num_workers=data.num_loading_workers,
        #     dataloader_drop_last=True,
        #     use_position_ids=isinstance(trainer.model_config, Qwen2Config),
        #     use_doc_masking=False,
        #     dataloader_pin_memory=True,
        # )
        # dist.barrier()

    else:
        raise ValueError(f"Unhandled case of `self.config.data.dataset`. Got: {data.dataset}")

    if sanity_check_dataloader_interval is not None:
        sanity_check_dataloader(
            dataloader,
            tokenizer_path=trainer.config.tokenizer.tokenizer_name_or_path,
            sanity_check_dataloader_interval=sanity_check_dataloader_interval,
        )

    return dataloader


def get_dataloader(
    trainer: DistributedTrainer, sanity_check_dataloader_interval: Optional[int] = None
) -> Dict[str, DataLoader]:
    dataloaders = {}

    # Print training plan
    log_rank("Training plan", logger=logger, level=logging.INFO, rank=0, is_separator=True)
    stages_info = "".join(
        f"[Stage {stage.name}] start from step {stage.start_training_step} \n" for stage in trainer.config.data_stages
    )
    full_log_message = f"There are {len(trainer.config.data_stages)} training stages \n{stages_info}"
    log_rank(full_log_message, logger=logger, level=logging.INFO, rank=0)

    current_stage = None
    # WARNING: we assume we train on last stage
    stage_idx = len(trainer.config.data_stages) - 1
    stage_args = trainer.config.data_stages[stage_idx]
    if trainer.iteration_step+1 == stage_args.start_training_step:
        log_rank(f"Starting new stage {stage_args.name}", logger=logger, level=logging.INFO, rank=0)
        # we start a new stage
        if stage_idx >= len(trainer.metadata.data_stages):
            trainer.metadata.data_stages.append(DataStageMetadata(
                name=stage_args.name,
                start_training_step=stage_args.start_training_step,
                consumed_train_samples=0,
                consumed_tokens_per_dataset_folder={},
                sequence_length=trainer.sequence_length,
            ))
    elif len(trainer.metadata.data_stages) < len(trainer.config.data_stages):
        raise ValueError(f"If you're trying to start a new stage, you need to set `start_training_step` to the step after the last stage's: {trainer.iteration_step+1}")
    current_stage = trainer.metadata.data_stages[stage_idx]
    cur_stage_consumed_train_samples = current_stage.consumed_train_samples
    consumed_tokens_per_dataset_folder = current_stage.consumed_tokens_per_dataset_folder
    stage_args_data = trainer.config.data_stages[stage_idx].data

    num_remaining_train_steps = compute_remain_train_steps_of_a_data_stage_from_ckp(
        current_stage, trainer.config, trainer.metadata
    ) # TODO: check this
    log_rank(
        f"Current stage: {current_stage.name} has {num_remaining_train_steps} remaining training steps and has consumed {cur_stage_consumed_train_samples} samples"
        f"Consumed tokens per dataset folder: {pformat(consumed_tokens_per_dataset_folder)}",
        logger=logger,
        level=logging.INFO,
        rank=0,
    )

    # warn that if seqlen of stage - 1 has changed, consumed_train_samples=0 so we'll assume we're reading from new folder (so that we can resume training)
    if current_stage.sequence_length != trainer.metadata.data_stages[-1].sequence_length:
        raise NotImplementedError("We don't support changing sequence length between stages yet")
        if current_stage.consumed_train_samples == 0:
            log_rank(
                f"Warning: The sequence length of the last stage has changed from {trainer.metadata.data_stages[-1].sequence_length} to {current_stage.sequence_length}. We'll assume we're reading from the beginning of the dataset folders.",
                logger=logger,
                level=logging.WARNING,
                rank=0,
            )
        else:
            # we're resuming training, so that's fine
            pass
        cur_stage_consumed_train_samples = current_stage.consumed_train_samples

    else:
        # Prepare last_stages_consumed_tokens_per_dataset_folder which will be used to offset BlendableDataset to avoid reseeing consumed tokens even when sampler has restarted for this stage
        last_stages_consumed_tokens_per_dataset_folder = {}
        for stage in trainer.metadata.data_stages[:-1]:
            for folder_path, consumed_tokens in stage.consumed_tokens_per_dataset_folder.items():
                last_stages_consumed_tokens_per_dataset_folder[folder_path] = last_stages_consumed_tokens_per_dataset_folder.get(folder_path, 0) + consumed_tokens  



    dataloaders[current_stage.name] = get_dataloader_from_data_stage(
        trainer,
        stage_args_data,
        consumed_train_samples_stage=cur_stage_consumed_train_samples,
        consumed_tokens_per_dataset_folder=consumed_tokens_per_dataset_folder,
        last_stages_consumed_tokens_per_dataset_folder=last_stages_consumed_tokens_per_dataset_folder,
        num_remaining_train_steps=num_remaining_train_steps,
        sanity_check_dataloader_interval=sanity_check_dataloader_interval,
    )
    return dataloaders


def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--config-file", type=str, required=True, help="Path to the YAML or python config file")
    parser.add_argument(
        "--sanity-check-dataloader-interval",
        type=int,
        default=None,
        help="Optional interval to print dataloader samples",
    )
    return parser.parse_args()


if __name__ == "__main__":
    args = get_args()
    config_file = args.config_file

    # Load trainer and data
    trainer = DistributedTrainer(config_file)
    dataloader = get_dataloader(trainer, args.sanity_check_dataloader_interval)

    # Train
    trainer.train(dataloader)