训练

NsysConfig

基类：BaseModel

nsys 性能分析的配置。

源代码位于 bionemo/llm/train.py

class NsysConfig(BaseModel):
    """Configuration for nsys profiling."""

    start_step: int = 0
    end_step: Optional[int] = None
    ranks: list[int] = field(default_factory=lambda: [0])

nemo_logger_factory(experiment_config, wandb_config)

基于提供的实验和 wandb 配置创建并返回 NeMoLogger 实例。

参数

名称	类型	描述	默认值
experiment_config	ExperimentConfig	包含实验设置（如结果目录、实验名称、检查点设置和记录器首选项）的配置对象。	必需
wandb_config	Optional[WandbConfig]	Weights and Biases 日志记录的可选配置对象。	必需

返回

类型	描述
NeMoLogger	nl.NeMoLogger：使用指定设置配置的 NeMoLogger 实例。

源代码位于 bionemo/llm/train.py

def nemo_logger_factory(experiment_config: ExperimentConfig, wandb_config: Optional[WandbConfig]) -> nl.NeMoLogger:
    """Creates and returns a NeMoLogger instance configured based on the provided experiment and wandb configurations.

    Args:
        experiment_config (ExperimentConfig): Configuration object containing experiment settings such as
            result directory, experiment name, checkpoint settings, and logger preferences.
        wandb_config (Optional[WandbConfig]): Optional configuration object for Weights and Biases logging.

    Returns:
        nl.NeMoLogger: An instance of NeMoLogger configured with the specified settings.
    """
    if experiment_config.create_checkpoint_callback:
        checkpoint_callback = nl_callbacks.ModelCheckpoint(
            save_last=experiment_config.save_last_checkpoint,
            monitor=experiment_config.metric_to_monitor_for_checkpoints,
            save_top_k=experiment_config.save_top_k,
            every_n_train_steps=experiment_config.save_every_n_steps,
            always_save_context=True,
            filename="{epoch}-{val_loss:.2f}-{step}-{consumed_samples}",  # Including step and consumed_samples in the checkpoint filename prevents duplicate filenames and bugs related to this.
        )
    else:
        checkpoint_callback = None

    nemo_logger = setup_nemo_lightning_logger(
        root_dir=experiment_config.result_dir,
        name=experiment_config.experiment_name,
        initialize_tensorboard_logger=experiment_config.create_tensorboard_logger,
        wandb_config=wandb_config,
        ckpt_callback=checkpoint_callback,
    )
    return nemo_logger

setup_trainer(parallel_config, training_config, callbacks=None, nsys_config=None)

使用指定的并行和训练配置来设置模型训练的训练器。

参数

名称	类型	描述	默认值
parallel_config	ParallelConfig	并行性配置，包括张量和流水线模型并行大小、设备数量和节点数量。	必需
training_config	TrainingConfig	训练配置，包括最大步数、加速器类型、验证批次限制、验证检查间隔和精度。	必需
callbacks	list	训练期间要使用的回调函数列表。默认为 None，在这种情况下，将使用默认回调（RichModelSummary 和 LearningRateMonitor）。	None
nsys_config	NsysConfig	nsys 性能分析的配置。如果为 None，则禁用。	None

返回

类型	描述
Trainer	nl.Trainer：配置好的训练器对象，可用于模型训练。

源代码位于 bionemo/llm/train.py

def setup_trainer(
    parallel_config: ParallelConfig,
    training_config: TrainingConfig,
    callbacks=None,
    nsys_config: NsysConfig | None = None,
) -> nl.Trainer:
    """Set up the trainer for model training using the specified parallel and training configurations.

    Args:
        parallel_config (ParallelConfig): Configuration for parallelism, including tensor and pipeline model parallel sizes,
                                          number of devices, and number of nodes.
        training_config (TrainingConfig): Configuration for training, including maximum steps, accelerator type,
                                          validation batch limit, validation check interval, and precision.
        callbacks (list, optional): List of callback functions to be used during training. Defaults to None,
                                    in which case default callbacks (RichModelSummary and LearningRateMonitor) are used.
        nsys_config (NsysConfig, optional): Configuration for nsys profiling. If None, is disabled.

    Returns:
        nl.Trainer: Configured trainer object ready for model training.
    """
    strategy = nl.MegatronStrategy(
        tensor_model_parallel_size=parallel_config.tensor_model_parallel_size,
        pipeline_model_parallel_size=parallel_config.pipeline_model_parallel_size,
        pipeline_dtype=get_autocast_dtype(training_config.precision),
        ddp=DistributedDataParallelConfig(
            check_for_nan_in_grad=True,
            overlap_grad_reduce=True,
            overlap_param_gather=False,  # TODO waiting for NeMo fix
            average_in_collective=True,
            use_distributed_optimizer=True,
        ),
        find_unused_parameters=True,
        gradient_as_bucket_view=True,
        ckpt_include_optimizer=True,
        ckpt_async_save=True,
        ckpt_parallel_load=True,
    )
    if callbacks is None:
        callbacks = [
            RichModelSummary(max_depth=4),
            LearningRateMonitor(),
        ]

    if training_config.include_perplexity:
        callbacks.append(PerplexityLoggingCallback())

    if training_config.gc_interval > 0:
        callbacks.append(
            nl_callbacks.GarbageCollectionCallback(
                gc_interval_train=training_config.gc_interval, gc_interval_val=training_config.gc_interval
            )
        )

    if nsys_config:
        if nsys_config.end_step is None:
            nsys_config.end_step = training_config.max_steps
        callbacks.append(
            nl_callbacks.NsysCallback(
                start_step=nsys_config.start_step,
                end_step=nsys_config.end_step,
                ranks=nsys_config.ranks,
                gen_shape=True,
            )
        )

    trainer = nl.Trainer(
        devices=parallel_config.num_devices,
        max_steps=training_config.max_steps,
        accelerator=training_config.accelerator,
        strategy=strategy,
        limit_val_batches=training_config.limit_val_batches,
        val_check_interval=training_config.val_check_interval,
        num_nodes=parallel_config.num_nodes,
        callbacks=callbacks,
        plugins=nl.MegatronMixedPrecision(
            precision=training_config.precision,
            params_dtype=get_autocast_dtype(training_config.precision),
            pipeline_dtype=get_autocast_dtype(training_config.precision),
            grad_reduce_in_fp32=False,
            autocast_enabled=False,
        ),
        enable_checkpointing=training_config.enable_checkpointing,
    )
    return trainer

train(bionemo_exposed_model_config, data_config, parallel_config, training_config, optim_config, experiment_config, wandb_config, nsys_config=None, resume_if_exists=True)

使用提供的配置训练 BioNemo 模型。使用 ExposedModelConfig 和 DataConfig 作为此方法的主要变体。

参数

名称	类型	描述	默认值
bionemo_exposed_model_config	ExposedModelConfig	公开的 BioNemo 模型的配置。	必需
data_config	DataConfig[DataModuleT]	数据模块的配置。	必需
parallel_config	ParallelConfig	并行训练的配置。	必需
training_config	TrainingConfig	训练参数的配置。	必需
optim_config	OptimizerSchedulerConfig	优化器和调度器的配置。	必需
experiment_config	ExperimentConfig	实验的配置。	必需
wandb_config	Optional[WandbConfig]	Weights and Biases 日志记录的配置。n	必需
nsys_config	Optional[NsysConfig]	nsys 性能分析的配置。如果为 None，则禁用。	None
resume_if_exists	bool	如果检查点存在，则恢复训练的标志。默认为 True。	True

源代码位于 bionemo/llm/train.py

def train(
    bionemo_exposed_model_config: ExposedModelConfig,
    data_config: DataConfig[DataModuleT],
    parallel_config: ParallelConfig,
    training_config: TrainingConfig,
    optim_config: OptimizerSchedulerConfig,
    experiment_config: ExperimentConfig,
    wandb_config: Optional[WandbConfig],
    nsys_config: Optional[NsysConfig] = None,
    resume_if_exists: bool = True,
):
    """Train a BioNemo model using the provided configurations. Uses the ExposedModelConfig and DataConfig as the primary variants for this method.

    Args:
        bionemo_exposed_model_config (ExposedModelConfig): Configuration for the exposed BioNemo model.
        data_config (DataConfig[DataModuleT]): Configuration for the data module.
        parallel_config (ParallelConfig): Configuration for parallel training.
        training_config (TrainingConfig): Configuration for training parameters.
        optim_config (OptimizerSchedulerConfig): Configuration for the optimizer and scheduler.
        experiment_config (ExperimentConfig): Configuration for the experiment.
        wandb_config (Optional[WandbConfig]): Configuration for Weights and Biases logging.n
        nsys_config (Optional[NsysConfig], optional): Configuration for nsys profiling. If None, is disabled.
        resume_if_exists (bool, optional): Flag to resume training if a checkpoint exists. Defaults to True.
    """
    bionemo_model_config = bionemo_exposed_model_config.exposed_to_internal_bionemo_model_config()
    pathlib.Path(data_config.result_dir).mkdir(parents=True, exist_ok=True)

    if experiment_config.save_every_n_steps != training_config.val_check_interval:
        logging.warning("Mutating training_config.save_every_n_steps to be equal to val_check_interval.")
        experiment_config.save_every_n_steps = training_config.val_check_interval

    global_batch_size = infer_global_batch_size(
        micro_batch_size=data_config.micro_batch_size,
        num_nodes=parallel_config.num_nodes,
        devices=parallel_config.num_devices,
        accumulate_grad_batches=parallel_config.accumulate_grad_batches,
        tensor_model_parallel_size=parallel_config.tensor_model_parallel_size,
        pipeline_model_parallel_size=parallel_config.pipeline_model_parallel_size,
    )

    data: DataModuleT = data_config.construct_data_module(global_batch_size)
    # TODO BioBertDataModule or BioBertTokenizer abstractions. We know all DataModuleT in this case has data.tokenizer,
    # although this constraint is not documented.

    # TODO: need an abstraction for LrSchedulerConfig
    if optim_config.lr_scheduler == "cosine":
        lr_scheduler = CosineAnnealingScheduler(
            max_steps=training_config.max_steps if optim_config.max_steps is None else optim_config.max_steps,
            min_lr=optim_config.lr / 100,
            warmup_steps=int(math.ceil(training_config.max_steps * optim_config.cosine_rampup_frac)),
            interval=optim_config.interval,
            monitor=optim_config.monitor,
            constant_steps=int(math.ceil(training_config.max_steps * optim_config.cosine_hold_frac)),
        )
    elif optim_config.lr_scheduler == "warmup_anneal":
        lr_scheduler = WarmupAnnealDecayHoldScheduler(
            warmup_steps=optim_config.warmup_steps,
            max_steps=training_config.max_steps if optim_config.max_steps is None else optim_config.max_steps,
            max_lr=optim_config.lr,
            min_lr=optim_config.lr / 10.0,
            anneal_percentage=0.10,
        )
    else:
        raise NotImplementedError(f"Scheduler {optim_config.lr_scheduler} not implemented.")

    optimizer = MegatronOptimizerModule(
        config=OptimizerConfig(
            lr=optim_config.lr,
            optimizer=optim_config.optimizer,
            use_distributed_optimizer=True,
            fp16=bionemo_model_config.fp16,
            bf16=bionemo_model_config.bf16,
        ),
        lr_scheduler=lr_scheduler,
    )

    model: BionemoLightningModule = biobert_lightning_module(
        config=bionemo_model_config, tokenizer=data.tokenizer, optimizer=optimizer
    )
    trainer: nl.Trainer = setup_trainer(parallel_config, training_config, nsys_config=nsys_config)
    nemo_logger: nl.NeMoLogger = nemo_logger_factory(experiment_config, wandb_config=wandb_config)

    llm.train(
        model=model,
        data=data,
        trainer=trainer,
        log=nemo_logger,
        resume=resume.AutoResume(
            resume_if_exists=resume_if_exists,
            resume_ignore_no_checkpoint=True,
        ),
    )