모델 병렬 처리를 이용한 모델 체크포인트 지정 및 미세 조정 - 아마존 SageMaker
분산 모델 체크포인트 지정분산 모델 미세 조정

기계 번역으로 제공되는 번역입니다. 제공된 번역과 원본 영어의 내용이 상충하는 경우에는 영어 버전이 우선합니다.

모델 병렬 처리를 이용한 모델 체크포인트 지정 및 미세 조정

SageMaker 모델 병렬화 라이브러리는 다양한 모델 병렬화 전략으로 분할된 모델 상태와 옵티마이저 상태를 저장하고, 학습을 다시 시작하고 미세 조정하려는 위치부터 연속 학습을 위한 체크포인트를 로드하는 체크포인트를 로드하는 체크포인트 API를 제공합니다. API는 모델 상태 및 최적화 프로그램 상태의 일부 또는 전체를 저장하는 옵션도 지원합니다.

분산 모델 체크포인트 지정

모델 병렬화 라이브러리의 PyTorch 프레임워크와 TensorFlow 사용 중인 모델 병렬화 라이브러리 버전에 따라 다음 주제 중 하나를 선택하십시오. SageMaker

분산 PyTorch 모델 체크포인트 지정 ( SageMaker 모델 병렬화 라이브러리 v1.10.0 이상용)

SageMaker 모델 병렬화 라이브러리는 분산 모델 상태 및 옵티마이저 상태의 전체 또는 일부 체크포인트를 저장하고 로드하기 위한 체크포인트 API를 제공합니다.

참고

모델 병렬화 라이브러리 v1.10.0 이상을 사용하는 경우 이 체크포인트 방법을 사용하는 PyTorch 것이 좋습니다. SageMaker

부분 체크포인트 지정

모델 병렬 처리로 훈련된 모델의 체크포인트를 저장하려면 부분 체크포인트 지정 옵션을 true(partial=True)로 설정한 상태에서 smdistributed.modelparallel.torch.save_checkpoint API를 사용하세요. 그러면 각 모델 파티션이 개별적으로 저장됩니다. 모델 및 최적화 프로그램 상태 외에, user_content 인수를 통해 사용자 지정 데이터를 추가로 저장할 수도 있습니다. 체크포인트가 지정된 모델, 최적화 프로그램 및 사용자 콘텐츠는 별도의 파일로 저장됩니다. save_checkpoint API 호출은 다음과 같은 구조로 체크포인트 폴더를 생성합니다.

- path
  - ${tag}_partial (folder for partial checkpoints)
    - model_rankinfo.pt
    - optimizer_rankinfo.pt
    - fp16_states_rankinfo.pt
    - user_content.pt
  - $tag (checkpoint file for full checkpoints)
  - user_content_$tag (user_content file for full checkpoints)
  - newest (a file that indicates the newest checkpoint)

부분 체크포인트에서 훈련을 재개하려면 partial=True(으)로 smdistributed.modelparallel.torch.resume_from_checkpoint API를 사용하고, 부분 체크포인트를 저장할 때 사용할 체크포인트 디렉터리 및 태그를 지정하세요. 참고로 모델 가중치의 실제 로드는 모델 파티셔닝 이후, 즉 smdistributed.modelparallel.torch.step 장식 훈련 단계 함수를 처음 실행할 때 발생합니다.

이 라이브러리는 부분 체크포인트를 저장할 경우 모델 파티션 결정을 파일 확장자가 .pt인 파일로 저장하기도 합니다. 반대로 부분 체크포인트에서 재개하면, 라이브러리가 파티션 결정 파일을 함께 로드합니다. 파티션 결정을 로드한 후에는 파티션을 변경할 수 없습니다.

다음 코드 스니펫은 교육 스크립트에서 체크포인트 API를 설정하는 방법을 보여줍니다. PyTorch 

import smdistributed.modelparallel.torch as smp

model = ...
model = smp.DistributedModel(model)
optimizer = ...
optimizer = smp.DistributedOptimizer(optimizer)
user_content = ...     # additional custom data
checkpoint_path = "/opt/ml/checkpoint/model_parallel"

# Save a checkpoint.
smp.save_checkpoint(
    path=checkpoint_path,
    tag=f"total_steps{total_steps}",
    partial=True,
    model=model,
    optimizer=optimizer,
    user_content=user_content
    num_kept_partial_checkpoints=5
)

# Load a checkpoint.
# This automatically loads the most recently saved checkpoint.
smp_checkpoint = smp.resume_from_checkpoint(
    path=checkpoint_path, 
    partial=True
)

전체 체크포인트 지정

추론 목적으로 최종 모델 아티팩트를 저장하려면 partial=False(으)로 smdistributed.modelparallel.torch.save_checkpoint API를 사용하세요. 이렇게 하면 모델 파티션이 서로 결합하여 단일 모델 아티팩트가 생성됩니다. 이렇게 해도 최적화 프로그램 상태는 서로 결합하지 않습니다.

전체 모델 체크포인트가 제공된 경우, 특정 가중치로 훈련을 초기화하기 위해 partial=False(으)로 smdistributed.modelparallel.torch.resume_from_checkpoint API를 사용할 수 있습니다. 이렇게 해도 최적화 프로그램 상태는 로드되지 않습니다.

참고

텐서 병렬 처리를 이용할 경우, state_dict은(는) 일반적으로 원본 모델 구현과 DistributedModel 구현 간에 번역되어야 합니다. 선택적으로 state_dict 번역 함수를 smdistributed.modelparallel.torch.resume_from_checkpoint의 인수로 제공할 수도 있습니다. 그러나 즉시 지원 모델의 경우에는 라이브러리가 이 번역을 자동으로 처리합니다.

다음 코드는 모델 병렬화로 학습된 모델을 완전히 체크포인팅하기 위해 체크포인트 API를 사용하는 방법의 예를 보여줍니다. PyTorch 

import smdistributed.modelparallel.torch as smp

model = ...
model = smp.DistributedModel(model)
optimizer = ...
optimizer = smp.DistributedOptimizer(optimizer)
user_content = ...     # additional custom data
checkpoint_path = "/opt/ml/checkpoint/model_parallel"

# Save a checkpoint.
smp.save_checkpoint(
    path=checkpoint_path,
    tag=f"total_steps{total_steps}",
    partial=False,
    model=model,
    optimizer=optimizer,
    user_content=user_content
    num_kept_partial_checkpoints=5
)

# Load a checkpoint.
# This automatically loads the most recently saved checkpoint.
smp_checkpoint = smp.resume_from_checkpoint(
    path=checkpoint_path, 
    partial=False
)

분산 PyTorch 모델 체크포인트 (v1.6.0과 v1.9.0 사이의 모델 병렬 처리 SageMaker 라이브러리용)

SageMaker 모델 병렬화 라이브러리는 텐서 병렬성을 사용하는 훈련 작업의 부분 또는 전체 체크포인트를 저장하는 Python 함수를 제공합니다. 다음 절차에서는 텐서 병렬 처리를 이용할 때 smp.save() 및 smp.load()을(를) 사용하여 체크포인트를 저장 및 로드하는 방법을 보여줍니다.

참고

및 v1.6.0과 v1.9.0 사이의 모델 병렬 처리 라이브러리를 사용하는 PyTorch 경우 이 체크포인트 방법을 사용하는 것이 좋습니다. 텐서 병렬 처리 SageMaker

  1. 모델 객체를 준비하고 라이브러리의 래퍼 함수 smp.DistributedModel()(으)로 해당 객체를 래핑하세요.

    model = MyModel(...)
model = smp.DistributedModel(model)

  2. 해당 모델을 위한 최적화 프로그램을 준비하세요. 모델 파라미터 세트는 최적화 프로그램 함수에 필요한 반복 가능 인수입니다. 모델 파라미터 세트를 준비하려면 각각의 모델 파라미터에 고유 ID를 할당하도록 model.parameters()을(를) 처리해야 합니다.

    모델 파라미터의 반복 가능 인수에 ID가 중복된 파라미터가 있으면 체크포인트가 지정된 최적화 프로그램 상태의 로드에 실패합니다. 최적화 프로그램의 고유 ID로 모델 파라미터의 반복 가능 인수를 생성하려면 다음을 참조하세요.

    unique_params = []
unique_params_set = set()
for p in model.parameters():
  if p not in unique_params_set:
    unique_params.append(p)
    unique_params_set.add(p)
del unique_params_set

optimizer = MyOpt(unique_params, ...)

  3. 라이브러리의 래퍼 함수 smp.DistributedOptimizer()을(를) 사용하여 최적화 프로그램을 래핑하세요.

    optimizer = smp.DistributedOptimizer(optimizer)

  4. smp.save()을(를) 사용하여 모델 및 최적화 프로그램 상태를 저장하세요. 체크포인트 저장 방식에 따라 다음 2가지 옵션 중 하나를 선택하세요.

    • 옵션 1: 단일 MP_GROUP의 각 mp_rank에 부분 모델을 저장하세요.

      model_dict = model.local_state_dict() # save a partial model
opt_dict = optimizer.local_state_dict() # save a partial optimizer state
# Save the dictionaries at rdp_rank 0 as a checkpoint
if smp.rdp_rank() == 0:
    smp.save(
        {"model_state_dict": model_dict, "optimizer_state_dict": opt_dict},
        f"/checkpoint.pt",
        partial=True,
    )

      텐서 병렬 처리를 이용할 경우, 라이브러리가 checkpoint.pt_{pp_rank}_{tp_rank} 형식으로 이름이 지정된 체크포인트 파일을 저장합니다.

      참고

      텐서 병렬 처리를 이용할 때는 if 명령문을 if smp.dp_rank() == 0 대신 if smp.rdp_rank() == 0(으)로 설정해야 합니다. 최적화 프로그램 상태가 텐서 병렬 처리로 샤딩되면 모든 감소 데이터 병렬 순위에 최적화 프로그램 상태의 자체 파티션이 저장되어야 합니다. 체크포인트에 잘못된 if 명령문을 사용하면 훈련 작업이 지연될 수 있습니다. 텐서 병렬화 if smp.dp_rank() == 0 없이 사용하는 방법에 대한 자세한 내용은 SageMaker Python SDK 문서의 저장 및 로드에 대한 일반 지침을 참조하십시오.

    • 옵션 2: 모델 전체를 저장하세요.

      if smp.rdp_rank() == 0:
    model_dict = model.state_dict(gather_to_rank0=True) # save the full model
    if smp.rank() == 0:
        smp.save(
            {"model_state_dict": model_dict},
            "/checkpoint.pt",
            partial=False,
        )
      참고

      전체 체크포인트를 지정하려면 다음 사항을 검토하세요.

      • gather_to_rank0=True을(를) 설정하면 0을(를) 제외한 모든 순위가 빈 사전을 반환합니다.

      • 전체 체크포인트를 지정할 경우에는 해당 모델에 대한 체크포인트 지정만 할 수 있습니다. 최적화 프로그램 상태에 대한 전체 체크포인트 지정은 현재 지원되지 않습니다.

      • 모델 전체를 smp.rank() == 0에 저장하기만 하면 됩니다.

  5. smp.load()를 사용하여 체크포인트를 로드하세요. 이전 단계에서의 체크포인트 지정 방식에 따라 다음 2가지 옵션 중 하나를 선택하세요.

    • 옵션 1: 부분 체크포인트를 로드하세요.

      checkpoint = smp.load("/checkpoint.pt", partial=True)
model.load_state_dict(checkpoint["model_state_dict"], same_partition_load=False)
optimizer.load_state_dict(checkpoint["optimizer_state_dict"])

      파티션이 변경되지 않는다는 것을 알면 로드가 더 빨리 이뤄지도록 model.load_state_dict()에서 same_partition_load=True을(를) 설정할 수 있습니다.

    • 옵션 2: 전체 체크포인트를 로드하세요.

      if smp.rdp_rank() == 0:
    checkpoint = smp.load("/checkpoint.pt", partial=False)
    model.load_state_dict(checkpoint["model_state_dict"])

      이 if smp.rdp_rank() == 0 조건은 필수 조건이 아니지만, 다양한 MP_GROUP 간에 중복되는 로드를 방지하는 데 도움이 될 수 있습니다. 전체 체크포인트 지정 최적화 프로그램 상태 사전은 현재 텐서 병렬 처리로 지원되지 않습니다.

분산 모델 체크포인트 TensorFlow

TensorFlow 모델 병렬화로 학습하는 동안 모델을 저장하려면 모델 병렬화 라이브러리에서 제공하는 다음 함수를 사용하십시오. SageMaker

분산 모델 미세 조정

훈련 스크립트에서는 미세 조정을 구성해야 합니다. 다음 코드 스니펫은 Hugging Face Transformers의 AutoModelForCausalLM 클래스를 사용하여 모듈 등록 smdistributed.model.parallel.torch 및 미세 조정 설정을 수정한 교육 스크립트의 예제 구조를 보여줍니다.

참고

smp.delayed_param_initialization 기능이 활성화된 분산 변환기(smp.DistributedModel()로 래핑된 변환기 모델)를 미세 조정하려면 FSx for Lustre 파일 시스템으로 미세 조정 작업을 구성해야 합니다. 지연된 파라미터 초기화 옵션으로 대규모 모델을 미세 조정하려는 경우, FSx for Lustre 파일 시스템을 설정해야 합니다.

import argparse
from transformers import AutoModelForCausalLM
import smdistributed.modelparallel
import smdistributed.modelparallel.torch as smp

def parse_args():

    parser = argparse.ArgumentParser()

    # set an arg group for model
    model_grp = parser.add_argument_group(
        title="model", description="arguments to describe model configuration"
    )

    ... # set up numerous args to parse from the configuration dictionary to the script for training

    # add arg for activating fine-tuning
    model_grp.add_argument(
        "--fine_tune",
        type=int,
        default=0,
        help="Fine-tune model from checkpoint or pretrained model",
    )

def main():
    """Main function to train GPT."""
    args = parse_args()

    ... # parse numerous args

    if args.fine_tune > 0 and args.delayed_param > 0 and smp.rank() == 0:
        pretrained_model = AutoModelForCausalLM.from_pretrained(
            args.model_name or args.model_dir
        )
        model_state_dict = pretrained_model.state_dict()
        path = os.path.join(args.model_dir, "fullmodel.pt")
        torch.save(model_state_dict, path)

    # create a Transformer model and wrap by smp.model_creation() 
    # with options to configure model parallelism parameters offered by SageMaker
    with smp.model_creation(
        tensor_parallelism=smp.tp_size() > 1 or args.use_distributed_transformer > 0,
        zero_init=args.use_distributed_transformer == 0,
        dtype=dtype,
        distribute_embedding=args.sharded_data_parallel_degree > 1 and smp.tp_size() > 1,
        use_alibi=args.alibi > 0,
        attention_in_fp32=args.attention_in_fp32 > 0,
        fp32_residual_addition=args.residual_addition_in_fp32 > 0,
        query_key_layer_scaling=args.query_key_layer_scaling > 0 and args.bf16 < 1,
        fused_softmax=args.fused_softmax > 0,
        fused_dropout=args.fused_dropout > 0,
        fused_bias_gelu=args.fused_bias_gelu > 0,
        flash_attention=args.flash_attention > 0,
    ):
        if args.fine_tune > 0 and args.delayed_param == 0:
            model = AutoModelForCausalLM.from_pretrained(
                args.model_name or args.model_dir
            )
        else:
            model = AutoModelForCausalLM.from_config(model_config)

    # wrap the model by smp.DistributedModel() to apply SageMaker model parallelism
    model = smp.DistributedModel(
        model, trace_device="gpu", backward_passes_per_step=args.gradient_accumulation
    )

    # wrap the optimizer by smp.DistributedOptimizer() to apply SageMaker model parallelism
    optimizer= ... # define an optimizer
    optimizer = smp.DistributedOptimizer(
        optimizer,
        static_loss_scale=None,
        dynamic_loss_scale=True,
        dynamic_loss_args={"scale_window": 1000, "min_scale": 1, "delayed_shift": 2},
    )

    # for fine-tuning, use smp.resume_from_checkpoint() to load a pre-trained model
    if args.fine_tune > 0 and args.delayed_param > 0:
        smp.resume_from_checkpoint(args.model_dir, tag="fullmodel.pt", partial=False)

트레이닝 스크립트와 Jupyter 노트북의 전체 예제를 보려면 예제 저장소의 GPT-2 예제를 참조하십시오. PyTorch SageMaker GitHub