Function initialize_model_and_tokenizer

initialize.py:55–116

(args)

def initialize_model_and_tokenizer(args):
    tokenizer = get_tokenizer(args)

    torch.distributed.barrier()
    start = time.time()

    for i in range(get_model_parallel_world_size()):
        if get_model_parallel_rank() == i:
            # Initialize model
            model = GLM130B(args).half()

            if args.from_quantized_checkpoint:
                assert args.quantization_bit_width is not None
                # Quantize model before moving to GPU
                model = quantize(model, args.quantization_bit_width)

            # Load checkpoint
            load_checkpoint(model, args)

            if args.quantization_bit_width is not None and not args.from_quantized_checkpoint:
                # Quantize model before moving to GPU
                model = quantize(model, args.quantization_bit_width)

            if args.bminf:
                import bminf

                if torch.distributed.get_rank() == 0:
                    print(f"> BMInf activated, memory limit: {args.bminf_memory_limit} GB")
                with torch.cuda.device(args.device):
                    model = bminf.wrapper(model, quantization=False, memory_limit=args.bminf_memory_limit << 30)
            else:
                model = model.to(args.device)
        if args.sequential_initialization:
            torch.distributed.barrier(group=get_model_parallel_group())

    torch.distributed.barrier()
    if torch.distributed.get_rank() == 0:
        print(f"> Model initialized in {time.time() - start:.1f}s")

    torch.cuda.empty_cache()
    model.eval()

    # generate rotary embedding cache
    original_parallel_output = model.transformer.parallel_output
    model.transformer.parallel_output = True
    with torch.no_grad():
        _, *_ = model(
            torch.ones(1, args.max_sequence_length, device=torch.cuda.current_device(), dtype=torch.int64),
            torch.arange(args.max_sequence_length, device=torch.cuda.current_device(), dtype=torch.int64).view(1, -1),
            torch.randn(
                1,
                1,
                args.max_sequence_length,
                args.max_sequence_length,
                device=torch.cuda.current_device(),
            )
            < 0.5,
        )