unitorch.cli.tasks

SupervisedTask

Tip

core/task/supervised is the section for configuration of SupervisedTask.

Standard supervised learning task with optional DDP, AMP, and EMA support.

Source code in src/unitorch/cli/tasks/supervised.py

def __init__(
    self,
    configure,
    model,
    datasets,
    local_rank: int = -1,  # GPU index for distributed training; -1 for single-GPU
    seed: int = 1123,  # global random seed for reproducibility
    cpu_offload: bool = False,  # keep model on CPU (e.g. for CPU-only environments)
):
    set_seed(seed)
    self.n_gpu = 1 if torch.cuda.is_available() else 0
    if dist.is_initialized():
        self.n_gpu = dist.get_world_size()

    self.config = configure
    self.model = model
    self.datasets = datasets
    self.local_rank = local_rank

    if self.local_rank != -1:
        torch.cuda.set_device(self.local_rank)

    if torch.cuda.is_available() and not cpu_offload:
        self.model = self.model.cuda()

    self.best_score = -np.inf

n_gpu instance-attribute

n_gpu = 1 if is_available() else 0

config instance-attribute

config = configure

model instance-attribute

model = model

datasets instance-attribute

datasets = datasets

local_rank instance-attribute

local_rank = local_rank

best_score instance-attribute

best_score = -inf

from_config classmethod

from_config(config, **kwargs)

Source code in src/unitorch/cli/tasks/supervised.py

@classmethod
@config_defaults_init("core/task/supervised")
def from_config(cls, config, **kwargs):
    try:
        torch.distributed.init_process_group(backend="nccl", init_method="env://")
    except Exception:
        logging.info("PyTorch is not in distributed mode")

    config.set_default_section("core/task/supervised")

    model = config.getoption("model", None)
    dataset = config.getoption("dataset", None)

    if model is not None:
        model = init_registered_module(model, config, registered_model)
    if dataset is not None:
        dataset = init_registered_module(dataset, config, registered_dataset)

    return dict(
        configure=config,
        model=model,
        datasets=dataset,
        local_rank=config.getdefault("core/cli", "local_rank", get_local_rank()),
        cpu_offload=config.getoption("cpu_offload", False),
    )

train

train(
    optim: str,
    loss_fn: str,
    score_fn: str,
    monitor_fns: Optional[Union[str, List[str]]] = None,
    scheduler: Optional[str] = None,
    from_ckpt_dir: str = "./from_ckpt",
    to_ckpt_dir: str = "./to_ckpt",
    train_batch_size: int = 128,
    dev_batch_size: int = 128,
    pin_memory: bool = True,
    num_workers: int = 4,
    save_optimizer: bool = True,
    save_scheduler: bool = True,
    save_checkpoint: str = "default",
    log_freq: int = 100,
    ckpt_freq: int = 10000,
    grad_acc_step: int = 1,
    max_grad_norm: float = 1.0,
    num_training_samples: int = 1000000000,
    epochs: int = 5,
    use_ema: bool = False,
    ema_decay: float = 0.9999,
    ema_tau: int = 2000,
    use_amp: bool = True,
)

Source code in src/unitorch/cli/tasks/supervised.py

@config_defaults_method("core/task/supervised")
def train(
    self,
    optim: str,  # registered optimizer name
    loss_fn: str,  # registered loss function name
    score_fn: str,  # registered scoring function name
    monitor_fns: Optional[
        Union[str, List[str]]
    ] = None,  # extra metrics logged at checkpoints
    scheduler: Optional[str] = None,  # registered LR scheduler name
    from_ckpt_dir: str = "./from_ckpt",  # directory to load pretrained weights from
    to_ckpt_dir: str = "./to_ckpt",  # directory to write checkpoints to
    train_batch_size: int = 128,  # per-GPU batch size for training
    dev_batch_size: int = 128,  # per-GPU batch size for validation
    pin_memory: bool = True,  # pin DataLoader memory for faster GPU transfer
    num_workers: int = 4,  # DataLoader worker processes
    save_optimizer: bool = True,  # include optimizer state in checkpoints
    save_scheduler: bool = True,  # include scheduler state in checkpoints
    save_checkpoint: str = "default",  # checkpoint policy: default/best/latest/every/all
    log_freq: int = 100,  # log training loss every N steps
    ckpt_freq: int = 10000,  # save checkpoint every N steps
    grad_acc_step: int = 1,  # gradient accumulation steps before optimizer update
    max_grad_norm: float = 1.0,  # gradient clipping max norm
    num_training_samples: int = 1_000_000_000,  # fallback total samples for iterable datasets
    epochs: int = 5,  # total training epochs
    use_ema: bool = False,  # maintain an EMA shadow model for evaluation
    ema_decay: float = 0.9999,  # EMA decay factor
    ema_tau: int = 2000,  # EMA warm-up steps
    use_amp: bool = True,  # enable automatic mixed precision (FP16)
):
    if self.local_rank in [-1, 0]:
        os.makedirs(to_ckpt_dir, exist_ok=True)

    if loss_fn is not None:
        loss_fn = init_registered_module(loss_fn, self.config, registered_loss)
    if score_fn is not None:
        score_fn = init_registered_module(score_fn, self.config, registered_score)
    if monitor_fns is not None:
        monitor_fns = [
            init_registered_module(fn, self.config, registered_score)
            for fn in monitor_fns
            if fn in registered_score
        ]

    if optim is not None and self.model is not None:
        optim = init_registered_module(
            optim,
            self.config,
            registered_optim,
            params=filter(lambda p: p.requires_grad, self.model.parameters()),
        )

    # Load pretrained weights, then resume from latest checkpoint if available
    if os.path.exists(from_ckpt_dir):
        self.model.from_checkpoint(from_ckpt_dir)
        optim.from_checkpoint(from_ckpt_dir, weight_name="pytorch_optim.bin")
    if os.path.exists(to_ckpt_dir):
        self.model.from_checkpoint(
            to_ckpt_dir, weight_name="pytorch_model_latest.bin"
        )
        optim.from_checkpoint(to_ckpt_dir, weight_name="pytorch_optim_latest.bin")

    info_path = os.path.join(to_ckpt_dir, "info.json")
    if os.path.exists(info_path):
        with open(info_path) as f:
            info = json.load(f)
    else:
        info = {}

    global_epoch = info.get("global_epoch", 0)
    global_step = info.get("global_step", 0)
    self.best_score = info.get("best_score", self.best_score)
    logging.info("best score so far: %s", self.best_score)

    self.ema_model = None
    if use_ema:
        self.ema_model = ExponentialMovingAverage(
            self.model,
            decay=ema_decay,
            tau=ema_tau,
            num_steps=info.get("num_ema_steps", 0),
        )
        if os.path.exists(from_ckpt_dir):
            self.ema_model.from_checkpoint(
                from_ckpt_dir, weight_name="pytorch_ema_model.bin"
            )
        if os.path.exists(to_ckpt_dir):
            self.ema_model.from_checkpoint(
                to_ckpt_dir, weight_name="pytorch_ema_model_latest.bin"
            )

    for name, param in self.model.named_parameters():
        logging.debug(
            "%s: trainable=%s dtype=%s shape=%s device=%s",
            name,
            param.requires_grad,
            param.dtype,
            param.shape,
            param.device,
        )

    global_rank = -1
    if self.n_gpu > 1:
        self.model = nn.parallel.DistributedDataParallel(
            self.model,
            device_ids=[self.local_rank],
            output_device=self.local_rank,
            find_unused_parameters=False,
            broadcast_buffers=False,
        )
        global_rank = dist.get_rank()

    train_sampler = DistributedSkipSampler if self.n_gpu > 1 else RandomSkipSampler
    dev_sampler = DistributedSampler if self.n_gpu > 1 else SequentialSampler

    dataset_train = self.datasets.get("train")
    dataset_dev = self.datasets.get("dev")

    iter_train = DataLoader(
        dataset_train,
        sampler=(
            train_sampler(dataset_train)
            if not isinstance(dataset_train, Iterable)
            else None
        ),
        batch_size=train_batch_size,
        shuffle=False,
        pin_memory=pin_memory,
        num_workers=num_workers,
        collate_fn=collate_fn,
    )
    iter_dev = DataLoader(
        dataset_dev,
        sampler=(
            dev_sampler(dataset_dev)
            if not isinstance(dataset_dev, Iterable)
            else None
        ),
        batch_size=dev_batch_size,
        shuffle=False,
        pin_memory=pin_memory,
        num_workers=num_workers,
        collate_fn=collate_fn,
    )

    if scheduler is not None:
        if not isinstance(dataset_train, Iterable):
            num_training_steps = int(
                epochs
                * len(dataset_train)
                // train_batch_size
                // max(1, self.n_gpu)
                // grad_acc_step
            )
        else:
            num_training_steps = int(
                epochs
                * num_training_samples
                // train_batch_size
                // max(1, self.n_gpu)
                // grad_acc_step
            )
        scheduler = init_registered_module(
            scheduler,
            self.config,
            registered_scheduler,
            optimizer=optim,
            num_training_steps=num_training_steps,
        )

    if scheduler is not None and os.path.exists(to_ckpt_dir):
        scheduler.from_checkpoint(
            to_ckpt_dir, weight_name="pytorch_scheduler_latest.bin"
        )

    # AMP gradient scaler; only created when use_amp=True
    scaler = torch.amp.GradScaler("cuda") if use_amp else None

    def _optimizer_step():
        """Unscale gradients (if AMP), clip, then step the optimizer."""
        if scaler is not None:
            scaler.unscale_(optim)
        nn.utils.clip_grad_norm_(self.model.parameters(), max_grad_norm)
        if scaler is not None:
            scaler.step(optim)
            scaler.update()
        else:
            optim.step()
        if scheduler is not None:
            scheduler.step()
        optim.zero_grad()
        if use_ema and self.ema_model is not None:
            base = self.model.module if self.n_gpu > 1 else self.model
            self.ema_model.step(base)

    def _snapshot(epoch, step):
        return save_snapshot(
            self.model.module if self.n_gpu > 1 else self.model,
            to_ckpt_dir,
            iter_dev,
            score_fn,
            monitor_fns,
            optim=optim if save_optimizer else None,
            scheduler=scheduler if save_scheduler else None,
            save_checkpoint=save_checkpoint,
            ema_model=self.ema_model if use_ema else None,
            best_score=self.best_score,
            info_path=info_path,
            local_rank=self.local_rank,
            global_epoch=epoch,
            global_step=step,
        )

    log_loss = 0.0
    dev_epoch = 0

    for e in range(epochs):
        torch.cuda.empty_cache()
        if e < global_epoch:
            continue

        if hasattr(dataset_train, "set_epoch"):
            dataset_train.set_epoch(e)
        if hasattr(dataset_train, "set_skip_step"):
            dataset_train.set_skip_step(global_step * train_batch_size)
        if hasattr(iter_train.sampler, "set_epoch"):
            iter_train.sampler.set_epoch(e)
        if hasattr(iter_train.sampler, "set_skip_step"):
            iter_train.sampler.set_skip_step(global_step * train_batch_size)

        self.model.train()
        is_update_step = True

        for step, (inputs, targets) in enumerate(iter_train):
            step = step + global_step
            is_update_step = False

            if torch.cuda.is_available():
                inputs = inputs.cuda()
                targets = targets.cuda()

            with torch.autocast(
                device_type="cuda" if torch.cuda.is_available() else "cpu",
                enabled=use_amp,
            ):
                outputs = self.model(**inputs.dict())
                loss = (
                    outputs.loss
                    if isinstance(outputs, LossOutputs)
                    else loss_fn(outputs=outputs, targets=targets)
                ) / grad_acc_step

            if scaler is not None:
                scaler.scale(loss).backward()
            else:
                loss.backward()

            log_loss += loss.item() * grad_acc_step

            if (step + 1) % grad_acc_step == 0:
                is_update_step = True
                _optimizer_step()

            if (step + 1) % log_freq == 0 and global_rank in [-1, 0]:
                avg_loss = log_loss / log_freq
                logging.info("epoch %d step %d: train/loss=%.6f", e, step, avg_loss)
                if wandb.is_available():
                    wandb.log({"epoch": e, "step": step, "train/loss": avg_loss})
                log_loss = 0.0

            if (step + 1) % ckpt_freq == 0:
                if hasattr(dataset_dev, "set_epoch"):
                    dataset_dev.set_epoch(dev_epoch)
                if hasattr(iter_dev.sampler, "set_epoch"):
                    iter_dev.sampler.set_epoch(dev_epoch)
                dev_epoch += 1
                self.best_score = _snapshot(e, step + 1)

        # Flush any remaining accumulated gradients at epoch end
        if not is_update_step:
            _optimizer_step()

        log_loss = 0.0

        if hasattr(dataset_dev, "set_epoch"):
            dataset_dev.set_epoch(dev_epoch)
        if hasattr(iter_dev.sampler, "set_epoch"):
            iter_dev.sampler.set_epoch(dev_epoch)
        dev_epoch += 1

        global_step = 0
        self.best_score = _snapshot(e + 1, 0)

eval

eval(
    monitor_fns: Union[str, List[str]],
    from_ckpt_dir: str = "./from_ckpt",
    dev_batch_size: int = 128,
    pin_memory: bool = True,
    num_workers: int = 4,
)

Source code in src/unitorch/cli/tasks/supervised.py

@torch.no_grad()
@config_defaults_method("core/task/supervised")
def eval(
    self,
    monitor_fns: Union[str, List[str]],  # list of registered scoring function names
    from_ckpt_dir: str = "./from_ckpt",  # directory to load model weights from
    dev_batch_size: int = 128,  # per-GPU batch size for evaluation
    pin_memory: bool = True,
    num_workers: int = 4,
):
    monitor_fns = [
        init_registered_module(fn, self.config, registered_score)
        for fn in monitor_fns
        if fn in registered_score
    ]

    if os.path.exists(from_ckpt_dir):
        self.model.from_checkpoint(from_ckpt_dir)

    global_rank = -1
    if self.n_gpu > 1:
        self.model = nn.parallel.DistributedDataParallel(
            self.model,
            device_ids=[self.local_rank],
            output_device=self.local_rank,
            find_unused_parameters=False,
            broadcast_buffers=False,
        )
        global_rank = dist.get_rank()

    dev_sampler = DistributedSampler if self.n_gpu > 1 else SequentialSampler
    dataset_dev = self.datasets.get("dev")
    iter_dev = DataLoader(
        dataset_dev,
        sampler=(
            dev_sampler(dataset_dev)
            if not isinstance(dataset_dev, Iterable)
            else None
        ),
        batch_size=dev_batch_size,
        shuffle=False,
        pin_memory=pin_memory,
        num_workers=num_workers,
        collate_fn=collate_fn,
    )

    results = infer(self.model.module if self.n_gpu > 1 else self.model, iter_dev)
    if global_rank in [-1, 0]:
        monitor(
            outputs=results.outputs,
            targets=results.targets,
            monitor_fns=monitor_fns,
        )

infer

infer(
    postprocess_fn: str,
    writer: str,
    test_batch_size: int = 128,
    pin_memory: bool = True,
    num_workers: int = 4,
    max_size: int = 10000,
    from_ckpt_dir: str = "./from_ckpt",
    output_header: Optional[List] = None,
    output_path: str = "./output.txt",
    postprocess_workers: int = 2,
)

Source code in src/unitorch/cli/tasks/supervised.py

@torch.no_grad()
@config_defaults_method("core/task/supervised")
def infer(
    self,
    postprocess_fn: str,  # registered postprocessing function name
    writer: str,  # registered writer name for output serialisation
    test_batch_size: int = 128,  # per-GPU batch size for inference
    pin_memory: bool = True,
    num_workers: int = 4,
    max_size: int = 10000,  # maximum queue depth for async postprocessing
    from_ckpt_dir: str = "./from_ckpt",  # directory to load model weights from
    output_header: Optional[
        List
    ] = None,  # column names to copy from raw dataset into output
    output_path: str = "./output.txt",  # file path for inference results
    postprocess_workers: int = 2,  # number of parallel postprocessing workers
):
    assert self.n_gpu <= 1, "inference only supports single-GPU mode"
    assert writer is not None

    output_dir = os.path.dirname(output_path)
    if output_dir:
        os.makedirs(output_dir, exist_ok=True)

    if postprocess_fn is not None:
        postprocess_fn = init_registered_process(postprocess_fn, self.config)

    writer = init_registered_module(
        writer, self.config, registered_writer, output_file=output_path
    )
    skip_step = writer.skip_n_samples

    if os.path.exists(from_ckpt_dir):
        self.model.from_checkpoint(from_ckpt_dir)

    sampler = SequentialSkipSampler if skip_step > 0 else SequentialSampler
    dataset_test = self.datasets.get("test")

    iter_test = DataLoader(
        dataset_test,
        sampler=(
            sampler(dataset_test)
            if not isinstance(dataset_test, Iterable)
            else None
        ),
        batch_size=test_batch_size,
        shuffle=False,
        pin_memory=pin_memory,
        num_workers=num_workers,
        collate_fn=collate_fn,
    )

    if skip_step > 0:
        if hasattr(dataset_test, "set_skip_step"):
            dataset_test.set_skip_step(skip_step)
        if hasattr(iter_test.sampler, "set_skip_step"):
            iter_test.sampler.set_skip_step(skip_step)

    # Build a parallel loader for raw dataset metadata (images, text) when available
    iter_data = None
    if hasattr(dataset_test, "dataset"):
        data_info = DatasetFeature(dataset_test.dataset)
        iter_data = DataLoader(
            deepcopy(data_info),
            sampler=(
                sampler(data_info)
                if not isinstance(dataset_test, Iterable)
                else None
            ),
            batch_size=test_batch_size,
            shuffle=False,
            pin_memory=pin_memory,
            num_workers=num_workers,
            collate_fn=None,
        )
        if skip_step > 0 and hasattr(iter_data.sampler, "set_skip_step"):
            iter_data.sampler.set_skip_step(skip_step)

    self.model.eval()
    start = time.time()

    data_queue = Queue(maxsize=max_size)
    msg_queue = Queue(maxsize=max_size)
    postprocess_list = [
        PostProcess(postprocess_fn, data_queue, msg_queue)
        for _ in range(postprocess_workers)
    ]
    for p in postprocess_list:
        p.start()

    io_process = IOProcess(msg_queue, writer=writer)
    io_process.start()

    if iter_data is None:
        for step, (inputs, _) in enumerate(iter_test):
            if torch.cuda.is_available():
                inputs = inputs.cuda()
            outputs = self.model(**inputs.dict()).cpu()
            data_queue.put((step, outputs))
    else:
        for step, ((inputs, _), raw_info) in enumerate(zip(iter_test, iter_data)):
            if torch.cuda.is_available():
                inputs = inputs.cuda()
            outputs = self.model(**inputs.dict()).cpu()
            if output_header is not None:
                raw_info = {k: raw_info[k] for k in output_header if k in raw_info}
                outputs.from_pandas(pd.DataFrame(raw_info))
            data_queue.put((step, outputs))

    data_queue.put((-1, GENERATE_FINISHED))
    for p in postprocess_list:
        p.join()

    msg_queue.put((-1, GENERATE_FINISHED))
    io_process.join()

    elapsed_ms = (time.time() - start) * 1000
    throughput = (len(dataset_test) - skip_step) / elapsed_ms * 1000
    logging.info("%.2f ms | %.2f samples/s", elapsed_ms, throughput)

DeepspeedTask

Tip

core/task/deepspeed/supervised is the section for configuration of DeepspeedTask.

Supervised learning task backed by DeepSpeed.

Source code in src/unitorch/cli/tasks/deepspeed.py

def __init__(
    self,
    configure,
    model,
    datasets,
    local_rank: int = -1,  # GPU index for distributed training; -1 for single-GPU
    seed: int = 1123,  # global random seed for reproducibility
    cpu_offload: bool = False,  # keep model on CPU (e.g. for ZeRO-Infinity offload)
):
    set_seed(seed)
    self.n_gpu = 1 if torch.cuda.is_available() else 0
    if dist.is_initialized():
        self.n_gpu = dist.get_world_size()

    self.config = configure
    self.model = model
    self.datasets = datasets
    self.local_rank = local_rank

    if self.local_rank != -1:
        torch.cuda.set_device(self.local_rank)

    if torch.cuda.is_available() and not cpu_offload:
        self.model = self.model.cuda()

    self.best_score = -np.inf

n_gpu instance-attribute

n_gpu = 1 if is_available() else 0

config instance-attribute

config = configure

model instance-attribute

model = model

datasets instance-attribute

datasets = datasets

local_rank instance-attribute

local_rank = local_rank

best_score instance-attribute

best_score = -inf

from_config classmethod

from_config(config, **kwargs)

Source code in src/unitorch/cli/tasks/deepspeed.py

@classmethod
@config_defaults_init("core/task/deepspeed/supervised")
def from_config(cls, config, **kwargs):
    try:
        deepspeed.init_distributed(dist_backend="nccl", init_method="env://")
    except Exception:
        logging.info("PyTorch is not in distributed mode")

    config.set_default_section("core/task/deepspeed/supervised")

    model = config.getoption("model", None)
    dataset = config.getoption("dataset", None)

    if model is not None:
        model = init_registered_module(model, config, registered_model)
    if dataset is not None:
        dataset = init_registered_module(dataset, config, registered_dataset)

    return dict(
        configure=config,
        model=model,
        datasets=dataset,
        local_rank=config.getdefault("core/cli", "local_rank", get_local_rank()),
        cpu_offload=config.getdefault(
            "core/task/deepspeed/supervised", "cpu_offload", False
        ),
    )

train

train(
    config_path: str,
    optim: str,
    loss_fn: str,
    score_fn: str,
    monitor_fns: Optional[Union[str, List[str]]] = None,
    from_ckpt_dir: str = "./from_ckpt",
    to_ckpt_dir: str = "./to_ckpt",
    train_batch_size: int = 128,
    dev_batch_size: int = 128,
    pin_memory: bool = True,
    num_workers: int = 4,
    save_optimizer: bool = False,
    save_scheduler: bool = False,
    save_checkpoint: str = "default",
    log_freq: int = 100,
    ckpt_freq: int = 10000,
    grad_acc_step: int = 1,
    learning_rate: Optional[float] = None,
    max_warmup_learning_rate: Optional[float] = None,
    num_warmup_steps: Optional[int] = None,
    epochs: int = 5,
    zero_stage: Optional[int] = None,
    merge_zero3_checkpoint: bool = True,
    exclude_freeze_parameters: bool = True,
    use_ema: bool = False,
    ema_decay: float = 0.9999,
    ema_tau: int = 2000,
)

Source code in src/unitorch/cli/tasks/deepspeed.py

@config_defaults_method("core/task/deepspeed/supervised")
def train(
    self,
    config_path: str,  # path to DeepSpeed JSON config file
    optim: str,  # registered optimizer name
    loss_fn: str,  # registered loss function name
    score_fn: str,  # registered scoring function name
    monitor_fns: Optional[
        Union[str, List[str]]
    ] = None,  # extra metrics logged at checkpoints
    from_ckpt_dir: str = "./from_ckpt",  # directory to load pretrained weights from
    to_ckpt_dir: str = "./to_ckpt",  # directory to write checkpoints to
    train_batch_size: int = 128,  # per-GPU micro-batch size for training
    dev_batch_size: int = 128,  # per-GPU batch size for validation
    pin_memory: bool = True,  # pin DataLoader memory for faster GPU transfer
    num_workers: int = 4,  # DataLoader worker processes
    save_optimizer: bool = False,  # include optimizer state in non-ZeRO-3 checkpoints
    save_scheduler: bool = False,  # include scheduler state in non-ZeRO-3 checkpoints
    save_checkpoint: str = "default",  # checkpoint policy: default/best/latest/every/all
    log_freq: int = 100,  # log training loss every N steps
    ckpt_freq: int = 10000,  # save checkpoint every N steps
    grad_acc_step: int = 1,  # gradient accumulation steps before optimizer update
    learning_rate: Optional[
        float
    ] = None,  # override optimizer lr in DeepSpeed config
    max_warmup_learning_rate: Optional[
        float
    ] = None,  # WarmupLR warmup_max_lr override
    num_warmup_steps: Optional[int] = None,  # WarmupLR warmup_num_steps override
    epochs: int = 5,  # total training epochs
    zero_stage: Optional[
        int
    ] = None,  # ZeRO optimisation stage (1/2/3); inferred from config if None
    merge_zero3_checkpoint: bool = True,  # merge ZeRO-3 shards into a single fp32 bin
    exclude_freeze_parameters: bool = True,  # skip frozen parameters when merging ZeRO-3
    use_ema: bool = False,  # maintain an EMA shadow model for evaluation
    ema_decay: float = 0.9999,  # EMA decay factor
    ema_tau: int = 2000,  # EMA warm-up steps
):
    if self.local_rank in [-1, 0]:
        os.makedirs(to_ckpt_dir, exist_ok=True)

    if loss_fn is not None:
        loss_fn = init_registered_module(loss_fn, self.config, registered_loss)
    if score_fn is not None:
        score_fn = init_registered_module(score_fn, self.config, registered_score)
    if monitor_fns is not None:
        monitor_fns = [
            init_registered_module(fn, self.config, registered_score)
            for fn in monitor_fns
            if fn in registered_score
        ]

    config_file = cached_path(config_path)
    with open(config_file) as f:
        config_dict = json.load(f)
    config_dict["train_micro_batch_size_per_gpu"] = train_batch_size

    if zero_stage is None:
        zero_stage = (
            nested_dict_value(config_dict, "zero_optimization", "stage") or 2
        )

    if os.path.exists(from_ckpt_dir):
        self.model.from_checkpoint(from_ckpt_dir)

    # Resume from latest checkpoint for non-ZeRO-3 (ZeRO-3 uses load_checkpoint below)
    if os.path.exists(to_ckpt_dir) and zero_stage != 3:
        self.model.from_checkpoint(
            to_ckpt_dir, weight_name="pytorch_model_latest.bin"
        )

    params = filter(lambda p: p.requires_grad, self.model.parameters())

    assert "optimizer" in config_dict
    update_nested_dict(config_dict, "zero_optimization", "stage", zero_stage)

    if learning_rate is not None:
        update_nested_dict(config_dict, "optimizer", "params", "lr", learning_rate)

    scheduler_type = nested_dict_value(config_dict, "scheduler", "type")
    if scheduler_type == "WarmupLR":
        if learning_rate is not None:
            update_nested_dict(
                config_dict, "scheduler", "params", "warmup_max_lr", learning_rate
            )
        if max_warmup_learning_rate is not None:
            update_nested_dict(
                config_dict,
                "scheduler",
                "params",
                "warmup_max_lr",
                max_warmup_learning_rate,
            )
        if num_warmup_steps is not None:
            update_nested_dict(
                config_dict,
                "scheduler",
                "params",
                "warmup_num_steps",
                num_warmup_steps,
            )

    info_path = os.path.join(to_ckpt_dir, "info.json")
    if os.path.exists(info_path):
        with open(info_path) as f:
            info = json.load(f)
    else:
        info = {}

    global_epoch = info.get("global_epoch", 0)
    global_step = info.get("global_step", 0)
    self.best_score = info.get("best_score", self.best_score)
    logging.info("best score so far: %s", self.best_score)

    self.ema_model = None
    if use_ema and zero_stage != 3:
        self.ema_model = ExponentialMovingAverage(
            self.model,
            decay=ema_decay,
            tau=ema_tau,
            num_steps=info.get("num_ema_steps", 0),
        )
        if os.path.exists(from_ckpt_dir):
            self.ema_model.from_checkpoint(
                from_ckpt_dir, weight_name="pytorch_ema_model.bin"
            )
        if os.path.exists(to_ckpt_dir):
            self.ema_model.from_checkpoint(
                to_ckpt_dir, weight_name="pytorch_ema_model_latest.bin"
            )

    for name, param in self.model.named_parameters():
        logging.debug(
            "%s: trainable=%s dtype=%s shape=%s device=%s",
            name,
            param.requires_grad,
            param.dtype,
            param.shape,
            param.device,
        )

    self.model, optim, _, scheduler = deepspeed.initialize(
        model=self.model,
        config=config_dict,
        model_parameters=params,
    )

    # ZeRO-3 checkpoints must be loaded after deepspeed.initialize
    if (
        os.path.exists(os.path.join(to_ckpt_dir, "pytorch_model_latest"))
        and zero_stage == 3
    ):
        self.model.load_checkpoint(
            os.path.join(to_ckpt_dir, "pytorch_model_latest")
        )

    global_rank = dist.get_rank() if self.n_gpu > 1 else -1
    train_sampler = DistributedSkipSampler if self.n_gpu > 1 else RandomSkipSampler
    dev_sampler = DistributedSampler if self.n_gpu > 1 else SequentialSampler

    dataset_train = self.datasets.get("train")
    dataset_dev = self.datasets.get("dev")

    iter_train = DataLoader(
        dataset_train,
        sampler=(
            train_sampler(dataset_train)
            if not isinstance(dataset_train, Iterable)
            else None
        ),
        batch_size=train_batch_size,
        shuffle=False,
        pin_memory=pin_memory,
        num_workers=num_workers,
        collate_fn=collate_fn,
    )
    iter_dev = DataLoader(
        dataset_dev,
        sampler=(
            dev_sampler(dataset_dev)
            if not isinstance(dataset_dev, Iterable)
            else None
        ),
        batch_size=dev_batch_size,
        shuffle=False,
        pin_memory=pin_memory,
        num_workers=num_workers,
        collate_fn=collate_fn,
    )

    snapshot_kwargs = dict(
        save_checkpoint=save_checkpoint,
        merge_checkpoint=merge_zero3_checkpoint,
        exclude_freeze_parameters=exclude_freeze_parameters,
        info_path=info_path,
        local_rank=self.local_rank,
    )

    def _snapshot(epoch, step):
        if zero_stage == 3:
            return save_snapshot_zero_3(
                self.model,
                to_ckpt_dir,
                iter_dev,
                score_fn,
                monitor_fns,
                best_score=self.best_score,
                global_epoch=epoch,
                global_step=step,
                **snapshot_kwargs,
            )
        return save_snapshot(
            self.model.module,
            to_ckpt_dir,
            iter_dev,
            score_fn,
            monitor_fns,
            optim=optim if save_optimizer else None,
            scheduler=scheduler if save_scheduler else None,
            ema_model=self.ema_model if use_ema else None,
            best_score=self.best_score,
            global_epoch=epoch,
            global_step=step,
            **snapshot_kwargs,
        )

    log_loss = 0.0
    dev_epoch = 0

    for e in range(epochs):
        torch.cuda.empty_cache()
        if e < global_epoch:
            continue

        if hasattr(dataset_train, "set_epoch"):
            dataset_train.set_epoch(e)
        if hasattr(dataset_train, "set_skip_step"):
            dataset_train.set_skip_step(global_step * train_batch_size)
        if hasattr(iter_train.sampler, "set_epoch"):
            iter_train.sampler.set_epoch(e)
        if hasattr(iter_train.sampler, "set_skip_step"):
            iter_train.sampler.set_skip_step(global_step * train_batch_size)

        self.model.train()
        is_update_step = True

        for step, (inputs, targets) in enumerate(iter_train):
            step = step + global_step
            is_update_step = False

            if torch.cuda.is_available():
                inputs = inputs.cuda()
                targets = targets.cuda()

            outputs = self.model(**inputs.dict())
            loss = (
                outputs.loss
                if isinstance(outputs, LossOutputs)
                else loss_fn(outputs=outputs, targets=targets)
            ) / grad_acc_step

            self.model.backward(loss)
            log_loss += loss.item() * grad_acc_step

            if (step + 1) % grad_acc_step == 0:
                is_update_step = True
                self.model.step()
                if use_ema and self.ema_model is not None:
                    self.ema_model.step(self.model.module)

            if (step + 1) % log_freq == 0 and global_rank in [-1, 0]:
                avg_loss = log_loss / log_freq
                logging.info("epoch %d step %d: train/loss=%.6f", e, step, avg_loss)
                if wandb.is_available():
                    wandb.log({"epoch": e, "step": step, "train/loss": avg_loss})
                log_loss = 0.0

            if (step + 1) % ckpt_freq == 0:
                if hasattr(dataset_dev, "set_epoch"):
                    dataset_dev.set_epoch(dev_epoch)
                if hasattr(iter_dev.sampler, "set_epoch"):
                    iter_dev.sampler.set_epoch(dev_epoch)
                dev_epoch += 1
                self.best_score = _snapshot(e, step + 1)

        # Flush remaining accumulated gradients at epoch end
        if not is_update_step:
            self.model.step()
            if use_ema and self.ema_model is not None:
                self.ema_model.step(self.model.module)

        log_loss = 0.0

        if hasattr(dataset_dev, "set_epoch"):
            dataset_dev.set_epoch(dev_epoch)
        if hasattr(iter_dev.sampler, "set_epoch"):
            iter_dev.sampler.set_epoch(dev_epoch)
        dev_epoch += 1

        global_step = 0
        self.best_score = _snapshot(e + 1, 0)

eval

eval(
    monitor_fns: Union[str, List[str]],
    from_ckpt_dir: str = "./from_ckpt",
    dev_batch_size: int = 128,
    pin_memory: bool = True,
    num_workers: int = 4,
)

Source code in src/unitorch/cli/tasks/deepspeed.py

@torch.no_grad()
@config_defaults_method("core/task/deepspeed/supervised")
def eval(
    self,
    monitor_fns: Union[str, List[str]],  # list of registered scoring function names
    from_ckpt_dir: str = "./from_ckpt",  # directory to load model weights from
    dev_batch_size: int = 128,  # per-GPU batch size for evaluation
    pin_memory: bool = True,
    num_workers: int = 4,
):
    monitor_fns = [
        init_registered_module(fn, self.config, registered_score)
        for fn in monitor_fns
        if fn in registered_score
    ]

    if os.path.exists(from_ckpt_dir):
        self.model.from_checkpoint(from_ckpt_dir)

    global_rank = -1
    if self.n_gpu > 1:
        self.model = nn.parallel.DistributedDataParallel(
            self.model,
            device_ids=[self.local_rank],
            output_device=self.local_rank,
            find_unused_parameters=False,
            broadcast_buffers=False,
        )
        global_rank = dist.get_rank()

    dev_sampler = DistributedSampler if self.n_gpu > 1 else SequentialSampler
    dataset_dev = self.datasets.get("dev")
    iter_dev = DataLoader(
        dataset_dev,
        sampler=(
            dev_sampler(dataset_dev)
            if not isinstance(dataset_dev, Iterable)
            else None
        ),
        batch_size=dev_batch_size,
        shuffle=False,
        pin_memory=pin_memory,
        num_workers=num_workers,
        collate_fn=collate_fn,
    )

    results = infer(self.model.module, iter_dev)
    if global_rank in [-1, 0]:
        monitor(
            outputs=results.outputs,
            targets=results.targets,
            monitor_fns=monitor_fns,
        )

infer

infer(
    postprocess_fn: str,
    writer: str,
    test_batch_size: int = 128,
    pin_memory: bool = True,
    num_workers: int = 4,
    max_size: int = 10000,
    from_ckpt_dir: str = "./from_ckpt",
    output_header: Optional[List] = None,
    output_path: str = "./output.txt",
    postprocess_workers: int = 2,
)

Source code in src/unitorch/cli/tasks/deepspeed.py

@torch.no_grad()
@config_defaults_method("core/task/deepspeed/supervised")
def infer(
    self,
    postprocess_fn: str,  # registered postprocessing function name
    writer: str,  # registered writer name for output serialisation
    test_batch_size: int = 128,  # per-GPU batch size for inference
    pin_memory: bool = True,
    num_workers: int = 4,
    max_size: int = 10000,  # maximum queue depth for async postprocessing
    from_ckpt_dir: str = "./from_ckpt",  # directory to load model weights from
    output_header: Optional[
        List
    ] = None,  # column names to copy from raw dataset into output
    output_path: str = "./output.txt",  # file path for inference results
    postprocess_workers: int = 2,  # number of parallel postprocessing workers
):
    assert self.n_gpu <= 1, "inference only supports single-GPU mode"
    assert writer is not None

    output_dir = os.path.dirname(output_path)
    if output_dir:
        os.makedirs(output_dir, exist_ok=True)

    if postprocess_fn is not None:
        postprocess_fn = init_registered_process(postprocess_fn, self.config)

    writer = init_registered_module(
        writer, self.config, registered_writer, output_file=output_path
    )
    skip_step = writer.skip_n_samples

    if os.path.exists(from_ckpt_dir):
        self.model.from_checkpoint(from_ckpt_dir)

    sampler = SequentialSkipSampler if skip_step > 0 else SequentialSampler
    dataset_test = self.datasets.get("test")

    iter_test = DataLoader(
        dataset_test,
        sampler=(
            sampler(dataset_test)
            if not isinstance(dataset_test, Iterable)
            else None
        ),
        batch_size=test_batch_size,
        shuffle=False,
        pin_memory=pin_memory,
        num_workers=num_workers,
        collate_fn=collate_fn,
    )

    if skip_step > 0:
        if hasattr(dataset_test, "set_skip_step"):
            dataset_test.set_skip_step(skip_step)
        if hasattr(iter_test.sampler, "set_skip_step"):
            iter_test.sampler.set_skip_step(skip_step)

    # Build a parallel loader for raw dataset metadata when available
    iter_data = None
    if hasattr(dataset_test, "dataset"):
        data_info = DatasetFeature(dataset_test.dataset)
        iter_data = DataLoader(
            deepcopy(data_info),
            sampler=(
                sampler(data_info)
                if not isinstance(dataset_test, Iterable)
                else None
            ),
            batch_size=test_batch_size,
            shuffle=False,
            pin_memory=pin_memory,
            num_workers=num_workers,
            collate_fn=None,
        )
        if skip_step > 0 and hasattr(iter_data.sampler, "set_skip_step"):
            iter_data.sampler.set_skip_step(skip_step)

    self.model.eval()
    start = time.time()

    data_queue = Queue(maxsize=max_size)
    msg_queue = Queue(maxsize=max_size)
    postprocess_list = [
        PostProcess(postprocess_fn, data_queue, msg_queue)
        for _ in range(postprocess_workers)
    ]
    for p in postprocess_list:
        p.start()

    io_process = IOProcess(msg_queue, writer=writer)
    io_process.start()

    if iter_data is None:
        for step, (inputs, _) in enumerate(iter_test):
            if torch.cuda.is_available():
                inputs = inputs.cuda()
            outputs = self.model(**inputs.dict()).cpu()
            data_queue.put((step, outputs))
    else:
        for step, ((inputs, _), raw_info) in enumerate(zip(iter_test, iter_data)):
            if torch.cuda.is_available():
                inputs = inputs.cuda()
            outputs = self.model(**inputs.dict()).cpu()
            if output_header is not None:
                raw_info = {k: raw_info[k] for k in output_header if k in raw_info}
                outputs.from_pandas(pd.DataFrame(raw_info))
            data_queue.put((step, outputs))

    data_queue.put((-1, GENERATE_FINISHED))
    for p in postprocess_list:
        p.join()

    msg_queue.put((-1, GENERATE_FINISHED))
    io_process.join()

    elapsed_ms = (time.time() - start) * 1000
    throughput = (len(dataset_test) - skip_step) / elapsed_ms * 1000
    logging.info("%.2f ms | %.2f samples/s", elapsed_ms, throughput)

MegatronTask

Tip

core/task/megatron/supervised is the section for configuration of MegatronTask.

Supervised learning task backed by Megatron-Core parallelism.

Source code in src/unitorch/cli/tasks/megatron.py

def __init__(
    self,
    configure,
    model,
    datasets,
    local_rank: int = -1,  # GPU index for distributed training; -1 for single-GPU
    seed: int = 1123,  # global random seed for reproducibility
):
    set_seed(seed)
    tensor_parallel.model_parallel_cuda_manual_seed(seed)
    self.config = configure
    self.model = model
    self.datasets = datasets
    self.local_rank = local_rank

    if self.local_rank != -1:
        torch.cuda.set_device(self.local_rank)

    if torch.cuda.is_available():
        self.model = self.model.cuda()

    self.best_score = -np.inf
    self.dp_rank = mpu.get_data_parallel_rank()
    self.dp_size = mpu.get_data_parallel_world_size()
    self.pp_rank = mpu.get_pipeline_model_parallel_rank()
    self.pp_size = mpu.get_pipeline_model_parallel_world_size()
    self.vp_size = mpu.get_virtual_pipeline_model_parallel_world_size()
    self.cp_rank = mpu.get_context_parallel_rank()
    self.cp_group = mpu.get_context_parallel_group()
    self.cp_size = mpu.get_context_parallel_world_size()
    self.tp_rank = mpu.get_tensor_model_parallel_rank()
    self.is_pp_last_rank = mpu.is_pipeline_last_stage(ignore_virtual=True)

config instance-attribute

config = configure

model instance-attribute

model = model

datasets instance-attribute

datasets = datasets

local_rank instance-attribute

local_rank = local_rank

best_score instance-attribute

best_score = -inf

dp_rank instance-attribute

dp_rank = get_data_parallel_rank()

dp_size instance-attribute

dp_size = get_data_parallel_world_size()

pp_rank instance-attribute

pp_rank = get_pipeline_model_parallel_rank()

pp_size instance-attribute

pp_size = get_pipeline_model_parallel_world_size()

vp_size instance-attribute

vp_size = get_virtual_pipeline_model_parallel_world_size()

cp_rank instance-attribute

cp_rank = get_context_parallel_rank()

cp_group instance-attribute

cp_group = get_context_parallel_group()

cp_size instance-attribute

cp_size = get_context_parallel_world_size()

tp_rank instance-attribute

tp_rank = get_tensor_model_parallel_rank()

is_pp_last_rank instance-attribute

is_pp_last_rank = is_pipeline_last_stage(
    ignore_virtual=True
)

_is_primary_rank property

_is_primary_rank: bool

True only on the single rank responsible for logging and checkpoint writes.

from_config classmethod

from_config(config, **kwargs)

Source code in src/unitorch/cli/tasks/megatron.py

@classmethod
@config_defaults_init("core/task/megatron/supervised")
def from_config(cls, config, **kwargs):
    try:
        torch.distributed.init_process_group(backend="nccl", init_method="env://")
    except Exception:
        logging.info("PyTorch is not in distributed mode")

    config.set_default_section("core/task/megatron/supervised")

    mpu.initialize_model_parallel(
        tensor_model_parallel_size=config.getoption(
            "tensor_model_parallel_size", 1
        ),
        pipeline_model_parallel_size=config.getoption(
            "pipeline_model_parallel_size", 1
        ),
        context_parallel_size=config.getoption("context_parallel_size", 1),
    )

    model = config.getoption("model", None)
    dataset = config.getoption("dataset", None)

    if model is not None:
        model = init_registered_module(model, config, registered_model)
    if dataset is not None:
        dataset = init_registered_module(dataset, config, registered_dataset)

    return dict(
        configure=config,
        model=model,
        datasets=dataset,
        local_rank=config.getdefault("core/cli", "local_rank", get_local_rank()),
    )

train

train(
    optim: str,
    loss_fn: str,
    score_fn: str,
    monitor_fns: Optional[Union[str, List[str]]] = None,
    scheduler: Optional[str] = None,
    from_ckpt_dir: str = "./from_ckpt",
    to_ckpt_dir: str = "./to_ckpt",
    train_batch_size: int = 128,
    dev_batch_size: int = 128,
    pin_memory: bool = True,
    num_workers: int = 4,
    log_freq: int = 100,
    ckpt_freq: int = 10000,
    grad_acc_step: int = 1,
    max_grad_norm: float = 1.0,
    num_training_samples: int = 1000000000,
    num_dev_samples: int = 10000,
    seq_length: Optional[int] = None,
    epochs: int = 5,
)

Source code in src/unitorch/cli/tasks/megatron.py

@config_defaults_method("core/task/megatron/supervised")
def train(
    self,
    optim: str,  # registered optimizer name
    loss_fn: str,  # registered loss function name
    score_fn: str,  # registered scoring function name
    monitor_fns: Optional[
        Union[str, List[str]]
    ] = None,  # extra metrics logged at checkpoints
    scheduler: Optional[str] = None,  # registered LR scheduler name
    from_ckpt_dir: str = "./from_ckpt",  # directory to load pretrained weights from
    to_ckpt_dir: str = "./to_ckpt",  # directory to write checkpoints to
    train_batch_size: int = 128,  # per-DP-rank micro-batch size for training
    dev_batch_size: int = 128,  # per-DP-rank batch size for validation
    pin_memory: bool = True,  # pin DataLoader memory for faster GPU transfer
    num_workers: int = 4,  # DataLoader worker processes
    log_freq: int = 100,  # log training loss every N steps
    ckpt_freq: int = 10000,  # save checkpoint every N steps
    grad_acc_step: int = 1,  # gradient accumulation microbatches
    max_grad_norm: float = 1.0,  # gradient clipping max norm
    num_training_samples: int = 1_000_000_000,  # fallback total samples for iterable datasets
    num_dev_samples: int = 10000,  # number of samples used per validation pass
    seq_length: Optional[
        int
    ] = None,  # sequence length hint passed to pipeline schedules
    epochs: int = 5,  # total training epochs
):
    if self.local_rank in [-1, 0]:
        os.makedirs(to_ckpt_dir, exist_ok=True)

    if loss_fn is not None:
        loss_fn = init_registered_module(loss_fn, self.config, registered_loss)
    if score_fn is not None:
        score_fn = init_registered_module(score_fn, self.config, registered_score)
    if monitor_fns is not None:
        monitor_fns = [
            init_registered_module(fn, self.config, registered_score)
            for fn in monitor_fns
            if fn in registered_score
        ]

    if optim is not None and self.model is not None:
        optim = init_registered_module(
            optim,
            self.config,
            registered_optim,
            params=filter(lambda p: p.requires_grad, self.model.parameters()),
        )

    if os.path.exists(from_ckpt_dir):
        self.model.from_checkpoint(from_ckpt_dir)
    if os.path.exists(os.path.join(to_ckpt_dir, "pytorch_model_latest")):
        self.model.from_checkpoint(
            os.path.join(to_ckpt_dir, "pytorch_model_latest")
        )

    info_path = os.path.join(to_ckpt_dir, "info.json")
    if os.path.exists(info_path):
        with open(info_path) as f:
            info = json.load(f)
    else:
        info = {}

    global_epoch = info.get("global_epoch", 0)
    global_step = info.get("global_step", 0)
    self.best_score = info.get("best_score", self.best_score)
    logging.info("best score so far: %s", self.best_score)

    for name, param in self.model.named_parameters():
        logging.debug(
            "%s: trainable=%s dtype=%s shape=%s device=%s",
            name,
            param.requires_grad,
            param.dtype,
            param.shape,
            param.device,
        )

    _ddp_config = DistributedDataParallelConfig(use_distributed_optimizer=False)
    if self.dp_size > 1:
        if isinstance(self.model, list):
            self.model = [
                DistributedDataParallel(
                    config=chunk.config,
                    module=chunk,
                    ddp_config=_ddp_config,
                )
                for chunk in self.model
            ]
        else:
            self.model = DistributedDataParallel(
                config=self.model.config,
                module=self.model,
                ddp_config=_ddp_config,
            )

    dataset_train = self.datasets.get("train")
    dataset_dev = self.datasets.get("dev")

    iter_train = iter(
        DataLoader(
            dataset_train,
            sampler=None,
            batch_size=train_batch_size,
            shuffle=False,
            pin_memory=pin_memory,
            num_workers=num_workers,
            collate_fn=collate_fn,
        )
    )
    iter_dev = iter(
        DataLoader(
            dataset_dev,
            sampler=None,
            batch_size=dev_batch_size,
            shuffle=False,
            pin_memory=pin_memory,
            num_workers=num_workers,
            collate_fn=collate_fn,
        )
    )

    if scheduler is not None:
        num_training_steps = int(
            epochs
            * num_training_samples
            // train_batch_size
            // max(1, self.dp_size)
            // grad_acc_step
        )
        scheduler = init_registered_module(
            scheduler,
            self.config,
            registered_scheduler,
            optimizer=optim,
            num_training_steps=num_training_steps,
        )

    # Select pipeline schedule based on parallelism configuration.
    # Interleaved VP requires model and data_iterator as lists (one per virtual chunk).
    if self.pp_size == 1:
        forward_backward_pipeline = forward_backward_no_pipelining
    elif self.vp_size is not None:
        forward_backward_pipeline = forward_backward_pipelining_with_interleaving
    else:
        forward_backward_pipeline = forward_backward_pipelining_without_interleaving

    def _make_pipeline_args(model, data_iterator):
        """Wrap model/data_iterator into lists when using interleaved VP schedule."""
        if self.vp_size is not None:
            if not isinstance(model, list):
                model = [model]
            if not isinstance(data_iterator, list):
                data_iterator = [data_iterator] * len(model)
        return model, data_iterator

    def _forward_step(data_iter, model):
        inputs, targets = next(data_iter)
        batch = _prepare_batch(inputs.dict(), targets.dict())
        outputs = model(**batch)
        if not isinstance(outputs, LossOutputs):
            # Non-loss output: return tensor + identity loss_func
            return outputs, lambda x: (x, x.detach())
        loss = outputs.loss
        # CP loss averaging is handled internally by Megatron's forward_step
        # (it multiplies by cp_world_size after dividing by num_microbatches).
        # Do NOT manually all_reduce or scale here to avoid double-scaling.
        return loss.unsqueeze(0), lambda x: (x.squeeze(0), x.detach().squeeze(0))

    num_train_steps = num_training_samples // train_batch_size // self.dp_size
    num_dev_steps = num_dev_samples // dev_batch_size // self.dp_size

    def _run_eval():
        self.model.eval()
        dev_loss = 0.0
        _model, _iter_dev = _make_pipeline_args(self.model, iter_dev)
        for _ in range(num_dev_steps):
            loss_list = forward_backward_pipeline(
                forward_step_func=_forward_step,
                data_iterator=_iter_dev,
                model=_model,
                num_microbatches=grad_acc_step,
                seq_length=seq_length,
                micro_batch_size=dev_batch_size,
                forward_only=True,
            )
            if self.is_pp_last_rank:
                dev_loss += loss_list[0].item() / num_dev_steps
        self.model.train()
        if self._is_primary_rank:
            logging.info("val/loss: %.6f", dev_loss)
            if wandb.is_available():
                wandb.log({"val/loss": dev_loss})
        return dev_loss

    def _save_checkpoint(epoch, step):
        base_model = getattr(self.model, "module", self.model)
        dev_loss = _run_eval()
        if -dev_loss > self.best_score:
            base_model.save_checkpoint(
                ckpt_dir=os.path.join(to_ckpt_dir, "pytorch_model"),
                weight_name="common.pt",
            )
            self.best_score = -dev_loss
        base_model.save_checkpoint(
            ckpt_dir=os.path.join(to_ckpt_dir, "pytorch_model_latest"),
            weight_name="common.pt",
        )
        info.update(
            best_score=self.best_score, global_epoch=epoch, global_step=step
        )
        if self._is_primary_rank:
            with open(info_path, "w") as f:
                json.dump(info, f, indent=4)

    log_loss = 0.0
    dev_epoch = 0

    for e in range(epochs):
        torch.cuda.empty_cache()
        if e < global_epoch:
            continue

        if hasattr(dataset_train, "set_epoch"):
            dataset_train.set_epoch(e)
        if hasattr(dataset_train, "set_skip_step"):
            dataset_train.set_skip_step(global_step)

        self.model.train()

        for step in range(num_train_steps):
            step = step + global_step

            _model, _iter_train = _make_pipeline_args(self.model, iter_train)
            loss_list = forward_backward_pipeline(
                forward_step_func=_forward_step,
                data_iterator=_iter_train,
                model=_model,
                num_microbatches=grad_acc_step,
                seq_length=seq_length,
                micro_batch_size=train_batch_size,
            )

            # Copy main_grad into .grad so standard optimizers can consume it
            for param in self.model.parameters():
                if hasattr(param, "main_grad") and param.main_grad is not None:
                    param.grad = param.main_grad

            if self.is_pp_last_rank:
                log_loss += loss_list[0].item()

            nn.utils.clip_grad_norm_(
                self.model.parameters(), max_norm=max_grad_norm
            )
            optim.step()
            if scheduler is not None:
                scheduler.step()
            optim.zero_grad()

            if (step + 1) % log_freq == 0 and self._is_primary_rank:
                avg_loss = log_loss / log_freq
                logging.info("epoch %d step %d: train/loss=%.6f", e, step, avg_loss)
                if wandb.is_available():
                    wandb.log({"epoch": e, "step": step, "train/loss": avg_loss})
                log_loss = 0.0

            if (step + 1) % ckpt_freq == 0:
                dist.barrier()
                if hasattr(dataset_dev, "set_epoch"):
                    dataset_dev.set_epoch(dev_epoch)
                dev_epoch += 1
                _save_checkpoint(e, step + 1)
                dist.barrier()

        log_loss = 0.0
        dist.barrier()

        if hasattr(dataset_dev, "set_epoch"):
            dataset_dev.set_epoch(dev_epoch)
        dev_epoch += 1
        _save_checkpoint(e + 1, 0)

        dist.barrier()
        global_step = 0