Function main

examples/tutorial/hybrid_parallel/train.py:42–136

()

def main():
    # launch from torch
    parser = colossalai.legacy.get_default_parser()
    args = parser.parse_args()
    colossalai.legacy.launch_from_torch(config=args.config)

    # get logger
    logger = get_dist_logger()
    logger.info("initialized distributed environment", ranks=[0])

    if hasattr(gpc.config, "LOG_PATH"):
        if gpc.get_global_rank() == 0:
            log_path = gpc.config.LOG_PATH
            if not os.path.exists(log_path):
                os.mkdir(log_path)
            logger.log_to_file(log_path)

    use_pipeline = is_using_pp()

    # create model
    model_kwargs = dict(
        img_size=gpc.config.IMG_SIZE,
        patch_size=gpc.config.PATCH_SIZE,
        hidden_size=gpc.config.HIDDEN_SIZE,
        depth=gpc.config.DEPTH,
        num_heads=gpc.config.NUM_HEADS,
        mlp_ratio=gpc.config.MLP_RATIO,
        num_classes=10,
        init_method="jax",
        checkpoint=gpc.config.CHECKPOINT,
    )

    if use_pipeline:
        pipelinable = PipelinableContext()
        with pipelinable:
            model = _create_vit_model(**model_kwargs)
        pipelinable.to_layer_list()
        pipelinable.policy = "uniform"
        model = pipelinable.partition(1, gpc.pipeline_parallel_size, gpc.get_local_rank(ParallelMode.PIPELINE))
    else:
        model = _create_vit_model(**model_kwargs)

    # count number of parameters
    total_numel = 0
    for p in model.parameters():
        total_numel += p.numel()
    if not gpc.is_initialized(ParallelMode.PIPELINE):
        pipeline_stage = 0
    else:
        pipeline_stage = gpc.get_local_rank(ParallelMode.PIPELINE)
    logger.info(f"number of parameters: {total_numel} on pipeline stage {pipeline_stage}")

    # use synthetic dataset
    # we train for 10 steps and eval for 5 steps per epoch
    train_dataloader = DummyDataloader(length=10, batch_size=gpc.config.BATCH_SIZE)
    test_dataloader = DummyDataloader(length=5, batch_size=gpc.config.BATCH_SIZE)

    # create loss function