Repository: oumi-ai/oumi
File: src/oumi/train.py (Python)

Code before:
        partial_trainer = functools.partial(
            create_trainer_fn,
            processing_class=tokenizer,
            args=config.training,
            train_dataset=dataset,
            eval_dataset=eval_dataset,
            **training_kwargs,
        )
        _verl_train(partial_trainer, checkpoint_location)
        return

    if is_distributed():
        init_distributed(timeout_minutes=config.training.nccl_default_timeout_minutes)

    # We support running FSDP Oumi training without being invoked from the Accelerate
    # launcher. We detect this with the following:
    # 1. Accelerate's environment variables aren't set
    # 2. We are running with a HF-family trainer (HF, TRL_SFT, TRL_DPO, TRL_GRPO)
    # 3. FSDP is enabled in the Oumi config
    # In this case, we mimic an Accelerate launcher run by setting the necessary
    # environment variables.
    # Note that normal Accelerate launcher runs won't be affected.
    if (
        not is_using_accelerate()
        and config.training.trainer_type != TrainerType.OUMI
        and config.fsdp.enable_fsdp
    ):
        accelerate_env_vars = prepare_accelerate_fsdp_run(config)
        logger.info(
            f"Set Accelerate environment variables for FSDP: {accelerate_env_vars}"
        )

    use_peft = config.training.use_peft and config.peft

    # Build model.
    model = build_model(
        model_params=config.model,
        peft_params=config.peft if use_peft else None,
        **(additional_model_kwargs or {}),
    )

    if use_peft:
        logger.info("Building PEFT model...")
        model = build_peft_model(
            model, config.training.enable_gradient_checkpointing, config.peft
        )

    if is_local_process_zero():
        log_number_of_model_parameters(model)
        if config.training.log_model_summary:
            log_model_summary(

Review comment (bug):
Nothing wrong with the changes as far as I can tell, but still it's a pretty big restructuring in `train()` function. Probably need to do some significant e2e testing that it doesn't lead to any unexpected side effects for misc models.

Code after:

        # We don't initialize the trainer here because it needs to run in a remote Ray
        # function.
        partial_trainer = functools.partial(
            create_trainer_fn,
            processing_class=tokenizer,
            args=config.training,
            train_dataset=dataset,
            eval_dataset=eval_dataset,
            **training_kwargs,
        )
        _verl_train(partial_trainer, checkpoint_location)
        return

    if is_distributed():
        init_distributed(timeout_minutes=config.training.nccl_default_timeout_minutes)

    # We support running FSDP Oumi training without being invoked from the Accelerate
    # launcher. We detect this with the following:
    # 1. Accelerate's environment variables aren't set
    # 2. We are running with a HF-family trainer (HF, TRL_SFT, TRL_DPO, TRL_GRPO)
    # 3. FSDP is enabled in the Oumi config
    # In this case, we mimic an Accelerate launcher run by setting the necessary
    # environment variables.
    # Note that normal Accelerate launcher runs won't be affected.
    if (
        not is_using_accelerate()
        and config.training.trainer_type != TrainerType.OUMI
        and config.fsdp.enable_fsdp
    ):
        accelerate_env_vars = prepare_accelerate_fsdp_run(config)
        logger.info(
            f"Set Accelerate environment variables for FSDP: {accelerate_env_vars}"
        )

    use_peft = config.training.use_peft and config.peft

    # Build model.
    model = build_model(
        model_params=config.model,
        peft_params=config.peft if use_peft else None,
        **(additional_model_kwargs or {}),
    )

    if use_peft:
        logger.info("Building PEFT model...")
        model = build_peft_model(
            model, config.training.enable_gradient_checkpointing, config.peft
        )

    if is_local_process_zero():