INNER CODE UNIT · Python

save_model_ddp

Guitaricet/relora · torchrun_main.py:192

def save_model_ddp(model, optimizer, scheduler, training_state_checkpoint, run_config, save_dir):
    global_rank = dist.get_rank()
    _time = time.time()

    if global_rank == 0:
        update_step = training_state_checkpoint["update_step"]
        os.makedirs(os.path.dirname(save_dir), exist_ok=True)

        _model = model.module
        _model.save_pretrained(save_dir)

    dist.barrier()
    if isinstance(optimizer, ZeroRedundancyOptimizer):
        logger.info("Started consolidating optimizer state dict")
        optimizer.consolidate_state_dict()
        logger.info(f"Consolidating optimizer state dict took {time.time() - _time:.2f} seconds")

    if global_rank == 0:

View source record →