[yongggg's] FSDP details and options

Getting Started with Fully Sharded Data Parallel(FSDP)

fsdp_auto_wrap_policy: FSDP에 fsdp_auto_wrap_policy를 적용하지 않는다면, FSDP는 전체 모델을 하나의 FSDP 단위에 배치하므로 계산과 메모리 효율성이 저하된다. 모델에 100개의 선형 레이어가 포함되어 있는 상황을 가정해보자. FSDP함수로 FSDP(model)을 수행하면, 전체 모델을 래핑하는 FSDP 단위가 하나만 있게 된다. 이 경우에 allgather step에서는 100개의 선형 레이어 모두에 대한 전체 parameter를 수집하므로 parameter sharding을 위해 CUDA 메모리를 절약하지 않는다. 또한 100개의 모든 선형 레이어에 대해 단 하나의 blocking allghater 호출만 있으므로 레이어 간의 통신 및 계산이 겹치지 않는다.

 

이를 방지하기 위해서, fsdp_auto_wrap_policy로 전달하면, 현재 FSDP 장치가 밀봉되고 지정된 조건(ex. size limit)이 충족되면 자동으로 새 장치가 자동으로 시작된다. 이렇게 하면 여러 개의 FSDP 장치가 있고 한 번에 한 개의 FSDP 장치만 collect하면 된다. 예를 들어, 5개의 FSDP 장치가 있고, 각 장치는 20개의 선형 레이어를 감싼다고 가정해보자. 그런 다음 forward에서 1st FSDP unit은 처음 20개의 선형 레이어에 대한 parameters를 수집하고 계산한 후 parameters를 폐기한 다음, 다음 20개의 선형 레이어로 이동한다. 따라서 모든 시점에서 각 Rank는 100개가 아닌 20개의 선형 레이어에 대한 parameters/grads만 구체화 한다.

def fsdp_main(rank, world_size, args):
    setup(rank, world_size)

    transform=transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
    ])

    dataset1 = datasets.MNIST('../data', train=True, download=True,
                        transform=transform)
    dataset2 = datasets.MNIST('../data', train=False,
                        transform=transform)

    sampler1 = DistributedSampler(dataset1, rank=rank, num_replicas=world_size, shuffle=True)
    sampler2 = DistributedSampler(dataset2, rank=rank, num_replicas=world_size)

    train_kwargs = {'batch_size': args.batch_size, 'sampler': sampler1}
    test_kwargs = {'batch_size': args.test_batch_size, 'sampler': sampler2}
    cuda_kwargs = {'num_workers': 2,
                    'pin_memory': True,
                    'shuffle': False}
    train_kwargs.update(cuda_kwargs)
    test_kwargs.update(cuda_kwargs)

    train_loader = torch.utils.data.DataLoader(dataset1,**train_kwargs)
    test_loader = torch.utils.data.DataLoader(dataset2, **test_kwargs)
    my_auto_wrap_policy = functools.partial(
        size_based_auto_wrap_policy, min_num_params=100
    )
    torch.cuda.set_device(rank)


    init_start_event = torch.cuda.Event(enable_timing=True)
    init_end_event = torch.cuda.Event(enable_timing=True)

    model = Net().to(rank)

    model = FSDP(model)

    optimizer = optim.Adadelta(model.parameters(), lr=args.lr)

    scheduler = StepLR(optimizer, step_size=1, gamma=args.gamma)
    init_start_event.record()
    for epoch in range(1, args.epochs + 1):
        train(args, model, rank, world_size, train_loader, optimizer, epoch, sampler=sampler1)
        test(model, rank, world_size, test_loader)
        scheduler.step()

    init_end_event.record()

    if rank == 0:
        print(f"CUDA event elapsed time: {init_start_event.elapsed_time(init_end_event) / 1000}sec")
        print(f"{model}")

    if args.save_model:
        # use a barrier to make sure training is done on all ranks
        dist.barrier()
        states = model.state_dict()
        if rank == 0:
            torch.save(states, "mnist_cnn.pt")

    cleanup()

이를 위해, 위으 코드에서 auto_wrap_policy를 정의하고 FSDP wrapper에 전달한다. 다음 예제에서 my_auto_wrap_policy는 이 계층의 parameters 수가 100보다 크면 layer가 FSDP에 의해 wrapping되거나 공유될 수 있다고 정의한다. 이 layer의 parameters 수가 100보다 작으면 FSDP에 의해 다른 작은 layer와 함께 wrapping된다. 최적의 자동 wrapping policy를 찾는 것이 어렵기 때문에 향후에 pytorch는 자동 튜닝을 추가할 것이다.

my_auto_wrap_policy = functools.partial(
        size_based_auto_wrap_policy, min_num_params=20000
    )
torch.cuda.set_device(rank)
model = Net().to(rank)

model = FSDP(model,
    fsdp_auto_wrap_policy=my_auto_wrap_policy)



cpu_offload: 모델이 매우 커서 FSDP를 사용해도 GPU에 맞지 않는 경우 cpu offload가 도움이 될 수 있다. 현재는 parameter 및 gradients의 cpu offload만 지원 된다. 이는 cpu_offload=CPUOFFload(offload_params=True)인자를 전달하여 활성화 할수 있다. 
이는 현재 parameters와 gradients가 최적화 프로그램과 함께 작동할 수 있는 동일한 장치에 있도록 CPU에 대한 그라데이션 offloading을 암시적으로 활성화한다. (이 API는 변경될 수 있다.)

이 기능을 사용하면, 호스트에서 장치로 tensor를 자주 복사하기 떄문에 훈련 속도가 상당히 느려질 수 있지만 메모리 효율성을 향상하고 더 큰 규모의 모델을 훈련하는 데 도움이 될 수 있다. 다음과 같이 FSDP wrapper에 추가하기만 하면 된다.

model = FSDP(model,
    fsdp_auto_wrap_policy=my_auto_wrap_policy,
    cpu_offload=CPUOffload(offload_params=True))

DDP보다 메모리 사용량이 낮은 것을 확인할 수 있다.

Advanced Model Training with Fully Sharded Data Parallel (FSDP)

Transformer Wrapping Policy

위에서 설명한 대로 auto_wrap_policy는 지정된 모델을 자동으로 sharding하고 모델, 최적화 프로그램 및 gradient shard를 별도의 FSDP 단위에 쉽게 배치할 수 있게 해주는 FSDP 기능 중 하나이다.

Transformer의 encoder-decoder와 같은 일부 아키텍처의 경우 embedding table과 같은 모델의 일부 부분이 encoder와 decoder 모두에 공유된다. 이 경우 encoder와 decoder 모두에 엑세스할 수 있도록 임베딩 table을 외부 FSDP 장치에 배치해야 한다. 또한, transformer에 대한 layer class를 등록함으로써 sharding 계획을 훨씬 더 효율적으로 통신할 수 있다. 

t5_auto_wrap_policy = functools.partial(
        transformer_auto_wrap_policy,
        transformer_layer_cls={
            T5Block,
        },
    )
torch.cuda.set_device(local_rank)


model = FSDP(model,
    fsdp_auto_wrap_policy=t5_auto_wrap_policy)

Mixed Precision

FSDP는 임의의 감소된 정밀도 유형(ex. fp16 or bfloat16)을 허용하는 유연한 mixed precision training을 지원한다. 현재 BFloat16은 Ampere GPU에서만 사용할 수 있으므로 사용하기 전에 기본 지원을 확인해야 한다. 예를 들어, v100에서는 BFloat16을 계속 실행할 수 있지만 기본적으로는 실행되지 않기 때문에 상당한 속도 저하가 발생할 수 있다. BFloat16이 기본적으로 지원되는지 확인하려면, 다음을 사용하면 된다.

bf16_ready = (
    torch.version.cuda
    and torch.cuda.is_bf16_supported()
    and LooseVersion(torch.version.cuda) >= "11.0"
    and dist.is_nccl_available()
    and nccl.version() >= (2, 10)
)

FSDP에서 혼합 정밀도의 장점 중 하나는 다음과 같이 parameters, gradients 및 버퍼에 대한 다양한 precision 수준에 대한 세부적인 제어를 제공한다는 것이다.

fpSixteen = MixedPrecision(
    param_dtype=torch.float16,
    # Gradient communication precision.
    reduce_dtype=torch.float16,
    # Buffer precision.
    buffer_dtype=torch.float16,
)

bfSixteen = MixedPrecision(
    param_dtype=torch.bfloat16,
    # Gradient communication precision.
    reduce_dtype=torch.bfloat16,
    # Buffer precision.
    buffer_dtype=torch.bfloat16,
)

fp32_policy = MixedPrecision(
    param_dtype=torch.float32,
    # Gradient communication precision.
    reduce_dtype=torch.float32,
    # Buffer precision.
    buffer_dtype=torch.float32,
)

Note that if a certain type (parameter, reduce, buffer)이 지정되지 않으면, 전혀 캐스팅되지 않는다.

이러한 유연성을 통해, 사용자는 gradients 통신을 reduced precision으로 설정하고 모든 parameters/buffer 계산을 full precision으로 수행하는 등의 세밀한 제어가 가능하다. 이는 노드 내 통신의 주요한 bottlenect 현상과 parameters/buffer로 인한 정확도 문제를 피하기 위해 full precision을 가져야 하는 경우에 잠재적으로 유용한다. 이는 다음 policy로 수행 가능하다.

grad_bf16 = MixedPrecision(reduce_dtype=torch.bfloat16)

 

Intializing FSDP Model on Device

FSDP는 device_id로 주어진 장치에서 입력 cpu 모듈을 초기화하는 device_id 인수를 지원한다. 이는 전체 모델이 single GPU에 맞지 않고 호스트의 cpu 메모리에 맞을 때 유용하다. device_id가 지정되면 FSDP는 cpu 기반 initialization보다 몇 배 빠르게 초기화하면서 GPU OOM 문제를 피하고 FSDP 단위로 모델을 지정된 장치로 이동시킨다.

torch.cuda.set_device(local_rank)

 model = FSDP(model,
        auto_wrap_policy=t5_auto_wrap_policy,
        mixed_precision=bfSixteen,
        device_id=torch.cuda.current_device())

Sharding Strategy

FSDP sharding strategy는 기본적으로 모델 parameters, gradients 및 optimizer state가 모든 layer에서 완전히 sharding되도록 설정한다(Zero3 sharding이라고도 함). optimizer state 및 gradients만 sharding되는 Zero2 sharding strategy에 관심이 있는 경우 FSDP는 다음과 같이 FSDP initialization에 "ShardingStrategy.FULL_SHARD" 대신에, "ShardingStrategy.SHARD_GRAD_OP"을 사용하여, sharding strategy를 전달함으로써 이 기능을 사용할 수 있다.

torch.cuda.set_device(local_rank)

 model = FSDP(model,
        auto_wrap_policy=t5_auto_wrap_policy,
        mixed_precision=bfSixteen,
        device_id=torch.cuda.current_device(),
        sharding_strategy=ShardingStrategy.SHARD_GRAD_OP # ZERO2)

이렇게 사용하면, FSDP의 통신 오버헤드가 줄어든다. 이 경우 forward 및 backward pass 이후에 전체 parameters를 유지한다. 다시 말해, backward pass 동안에 all_gather를 절약할 수 있으므로 메모리 사용량은 늘지만, 대신 통신을 덜 할 수 있다. full model params는 backwards의 끝 부분에서 해제되고 all_gather는 다음 forward pass에서 발생한다.

Backward Prefetch

backward prefetch setting은 next FSDP unit의 parameters를 요청해야 하는 시점을 제어한다. BACKWARD_PRE로 설정하면, 다음 FSDP의 unit params가 요청되기 시작하여 현재 unit의 계산이 시작되기 전에 더 일찍 도착할 수 있다. 이는 all_gather 통신 및 gradient computation과 중복되며, 약간 더 높은 memory 소비를 대가로 훈련 속도를 높일 수 있다. 다음과 같은 FSDP wrapper에 사용할 수 있다.

torch.cuda.set_device(local_rank)

 model = FSDP(model,
        auto_wrap_policy=t5_auto_wrap_policy,
        mixed_precision=bfSixteen,
        device_id=torch.cuda.current_device(),
        backward_prefetch = BackwardPrefetch.BACKWARD_PRE)

backward_prefetch에는 BACKWARD_PRE 및 BACKWARD_POST 두 가지 모드가 있다. BACKWARD_POST는 현재 FSDP unit 처리가 완료될 때까지 다음 FSDP unit의 params를 요청하지 않으므로 memory overhead를 최소화할 수 있다. 경우에 따라 BACKWARD_PRE를 사용하면 모델 학습 속도를 최대 2-10%까지 높일 수 있으며, 더 큰 모델에서는 훨씬 더 빠른 속도향상이 가능하다.

Model Checkpoint Saving, by streaming to the Rank0 CPU

local model과 동일한 방식으로 모델을 저장하는 FULL_STATE_DICT save를 사용하여 model checkpoint를 저장하기 위해 Pytorch 1.12 version에서는 더 큰 모델의 저장을 지원하는 몇 가지 utilities를 제공한다. 

먼저, FullStateDictConfig를 지정하여 state_dict를 rank 0에만 채우고 cpu로 offload할 수 있다.

이 구성을 사용하면, FSDP는 모두 모델 parameters를 모아서 rank 0에서만 cpu로 offload한다. state_dict가 최종적으로 저장되면 state_dict는 rank 0에만 채워지고 cpu tensor tensors를 포함한다. 따라서 single GPU memory 보다 큰 모델의 경우 잠재적인 OOM을 피할 수 있으며, 크기가 대략 사용자의 컴퓨터에서 사용 가능한 CPU RAM인 checkpoint model을 사용할 수 있다.

사용법은 다음과 같다.

save_policy = FullStateDictConfig(offload_to_cpu=True, rank0_only=True)
with FSDP.state_dict_type(
            model, StateDictType.FULL_STATE_DICT, save_policy
        ):
            cpu_state = model.state_dict()
if rank == 0:
 save_name = file_save_name + "-" + time_of_run + "-" + currEpoch
 torch.save(cpu_state, save_name)