maxtext
v1.0.0
MaxText는 순수 Python/Jax로 작성되었으며 학습 및 추론을 위해 Google Cloud TPU 및 GPU를 대상으로 하는 고성능 , 확장성이 뛰어난 오픈 소스 LLM입니다. MaxText는 Jax 및 XLA 컴파일러의 강력한 덕분에 단순함과 "최적화 없는" 상태를 유지하면서 높은 MFU를 달성하고 단일 호스트에서 매우 큰 클러스터로 확장됩니다.
MaxText는 연구 및 생산 분야 모두에서 야심찬 LLM 프로젝트의 출발점이 되는 것을 목표로 합니다. 우리는 사용자가 즉시 MaxText를 실험한 후 필요에 맞게 MaxText를 분기하고 수정하는 것으로 시작하는 것을 권장합니다.
우리는 MaxText를 사용하여 int8에서 고성능의 잘 수렴되는 훈련을 시연하고 훈련을 ~51K 칩으로 확장했습니다.
주요 지원 기능:
MaxText를 처음 실행하는 경우 구체적인 지침을 제공합니다.
MaxText는 다양한 개방형 모델의 훈련 및 추론을 지원합니다. 자세한 내용은 시작하기 폴더의 사용자 가이드를 따르세요.
추가로 유용한 가이드:
시작 가이드 외에도 지속적으로 추가되는 다른 MaxText 기능이 항상 있습니다! 전체 엔드투엔드 테스트 모음은 end_to_end에 있습니다. 우리는 밤마다 리듬을 타서 운영합니다. MaxText를 이해하는 데 좋은 소스가 될 수 있습니다. 또는 거의 지속적으로 실행되는 연속 단위 테스트를 볼 수도 있습니다.
이러한 결과 재현에 대한 자세한 내용은 MaxText/configs/README.md에서 확인할 수 있습니다.
| 매개변수 수 | 가속기 유형 | TFLOP/칩/초 | 모델 플롭 활용도(MFU) |
|---|---|---|---|
| 32B | v5p-128 | 3.28e+02 | 71.47% |
| 64B | v5p-128 | 3.23e+02 | 70.31% |
| 128B | v5p-256 | 3.15e+02 | 68.68% |
| 128B | v5p-512 | 3.15e+02 | 68.53% |
| 256B | v5p-1024 | 3.16e+02 | 68.82% |
| 512B | v5p-1024 | 2.94e+02 | 63.99% |
| 1024B | v5p-2048 | 2.49e+02 | 64.05% |
| 1024B | v5p-4096 | 2.97e+02 | 64.80% |
| 1160B | v5p-7680 | 2.95e+02 | 64.27% |
| 1160B | v5p-12288 | 3.04e+02 | 66.23% |
16B, 32B, 64B 및 128B 모델용. MaxText/configs/v5e/의 전체 실행 구성을 16b.sh , 32b.sh , 64b.sh , 128b.sh 로 확인하세요.
| 하드웨어 | 16B TFLOP/초/칩 | 16B MFU | 32B TFLOP/초/칩 | 320억 MFU | 64B TFLOP/초/칩 | 64BMFU | 128B TFLOP/초/칩 | 128B MFU |
|---|---|---|---|---|---|---|---|---|
| 1x v5e-256 | 120 | 61.10% | 132 | 66.86% | 118 | 59.90% | 110 | 56.06% |
| 2x v5e-256 | 117 | 59.37% | 128 | 64.81% | 112 | 56.66% | 110 | 55.82% |
| 4x v5e-256 | 117 | 59.14% | 126 | 64.10% | 110 | 55.85% | 108 | 54.93% |
| 8x v5e-256 | 115 | 58.27% | 125 | 63.67% | 108 | 54.96% | 104 | 52.93% |
| 16x v5e-256 | 111 | 56.56% | 123 | 62.26% | 105 | 53.29% | 100 | 50.86% |
| 32x v5e-256 | 108 | 54.65% | 119 | 60.40% | 99 | 50.18% | 91 | 46.25% |
MaxText는 PyTorch로 작성되고 Nvidia GPU를 대상으로 하는 우아한 독립형 GPT 구현인 MinGPT/NanoGPT에서 많은 영감을 받았습니다. MaxText는 더 복잡하여 더 많은 산업 표준 모델을 지원하고 수만 개의 칩으로 확장됩니다. 궁극적으로 MaxText는 해당 코드베이스로 가장 최근에 보고된 17%의 3배가 넘는 MFU를 가지며 대규모 확장이 가능하고 효율적인 자동 회귀 디코딩을 위한 키-값 캐시를 구현합니다.
MaxText는 Nvidia GPU를 대상으로 매우 잘 조정된 LLM 구현인 Nvidia/Megatron-LM과 더 유사합니다. 두 가지 구현은 비슷한 MFU를 달성합니다. 코드베이스의 차이는 다양한 프로그래밍 전략을 강조합니다. MaxText는 순수한 Python으로 XLA 컴파일러에 크게 의존하여 고성능을 달성합니다. 이와 대조적으로 Megatron-LM은 Python과 CUDA가 혼합되어 있으며 잘 최적화된 CUDA 커널을 사용하여 고성능을 달성합니다.
MaxText는 Pax와도 비슷합니다. Pax와 마찬가지로 MaxText는 Jax에서 LLM의 확장 가능한 고성능 구현을 제공합니다. Pax는 강력한 구성 매개변수를 활성화하여 개발자가 구성 매개변수를 편집하여 모델을 변경할 수 있도록 하는 데 중점을 둡니다. 대조적으로 MaxText는 사용자가 소스 코드를 분기하고 직접 편집하여 확장하도록 권장하는 다양한 LLM의 간단하고 구체적인 구현입니다.
가속기에서 SPMD(단일 프로그램, 다중 데이터) 작업을 실행할 때 오류가 발생하거나 어떤 이유로 VM이 정지/충돌하는 경우 전체 프로세스가 정지될 수 있습니다. 이 시나리오에서 스택 추적을 캡처하면 TPU VM에서 실행되는 작업의 문제를 식별하고 해결하는 데 도움이 됩니다.
다음 구성은 스택 추적을 수집하여 오류를 디버깅하거나 프로그램이 어딘가에 멈추거나 중단된 경우 도움이 됩니다. MaxText/configs/base.yml 에서 매개변수 값을 적절하게 변경하십시오.
collect_stack_trace: True 설정합니다. 이 설정은 디버깅에 도움이 되도록 프로그램에 대한 추적을 주기적으로 덤프합니다. 이를 비활성화하려면 collect_stack_trace: False 설정하세요.stack_trace_to_cloud: False 설정합니다. stack_trace_to_cloud: True TPU의 /tmp/debugging 에 임시 파일을 생성하여 스택 추적을 저장합니다. 임시 디렉터리에서 gcp 프로젝트의 클라우드 로깅으로 추적을 주기적으로 업로드하는 TPU VM에서 실행되는 에이전트가 있습니다. 다음 쿼리를 사용하여 Cloud Logging의 로그 탐색기에서 추적을 볼 수 있습니다. logName="projects/<project_name>/logs/tpu.googleapis.com%2Fruntime_monitor"
jsonPayload.verb="stacktraceanalyzer"
stack_trace_interval_seconds 각 스택 추적 수집 이벤트 사이의 기간(초)을 나타냅니다. stack_trace_interval_seconds: 600 설정하면 600초(10분)마다 스택 추적을 수집합니다.관련 PyPI 패키지는 다음과 같습니다: https://pypi.org/project/cloud-tpu-diagnostics.
훈련 실행을 미리 컴파일하기 위해 train_compile.py 도구를 제공합니다. 이 도구를 사용하면 전체 클러스터를 사용하지 않고도 대상 하드웨어(예: 다수의 v5e 장치)에 대해 train.py 에서 기본 train_step 컴파일할 수 있습니다.
TPU 클러스터용으로 사전 컴파일하려면 CPU 또는 다른 제품군의 단일 VM만 사용할 수 있습니다. 이 편집은 두 가지 주요 목표에 도움이 됩니다.
per_device_batch_size 가 너무 높게 설정된 경우와 같은 메모리 부족(OOM) 정보에 플래그를 지정하고 대상 하드웨어에서 컴파일된 것처럼 동일한 OOM 스택 추적을 사용합니다.
사전 컴파일을 저장한 다음 대상 하드웨어에서 빠른 시작 및 재시작 시간을 위해 로드할 수 있습니다.
train_compile.py 도구는 train.py 와 긴밀하게 연결되어 있으며 동일한 구성 파일 configs/base.yml 사용합니다. TPU에서 실행할 필요는 없지만 다른 종속 항목 외에 jax[tpu] 설치해야 하므로 아직 설치하지 않은 경우 setup.sh 실행하여 설치하는 것이 좋습니다.
위에 나열된 종속성을 설치하고 나면 미리 컴파일할 준비가 된 것입니다.
# Run the below on a single machine, e.g. a CPU
python3 MaxText/train_compile.py MaxText/configs/base.yml compile_topology=v5e-256 compile_topology_num_slices=2
global_parameter_scale=16 per_device_batch_size=4
그러면 2개의 v5e 포드에서 16B 매개변수 MaxText 모델이 컴파일됩니다.
다음은 컴파일된 train_step 을 저장한 다음 로드하는 예입니다. 저장부터 시작합니다.
1단계: AOT 실행 및 컴파일된 함수 저장
# Run the below on a single machine, e.g. a CPU
export LIBTPU_INIT_ARGS="--xla_enable_async_all_gather=true"
python3 MaxText/train_compile.py MaxText/configs/base.yml compile_topology=v5e-256
compile_topology_num_slices=2
compiled_trainstep_file=my_compiled_train.pickle global_parameter_scale=16
per_device_batch_size=4 steps=10000 learning_rate=1e-3
2단계: train.py를 실행하고 컴파일된 함수 로드
컴파일된 train_step을 로드하려면 compiled_trainstep_file=my_compiled_train.pickle train.py 에 전달하면 됩니다.
# Run the below on each host of the target hardware, e.g. each host on 2 slices of v5e-256
export LIBTPU_INIT_ARGS="--xla_enable_async_all_gather=true"
python3 MaxText/train.py MaxText/configs/base.yml run_name=example_load_compile
compiled_trainstep_file=my_compiled_train.pickle
global_parameter_scale=16 per_device_batch_size=4 steps=10000 learning_rate=1e-3
base_output_directory=gs://my-output-bucket dataset_path=gs://my-dataset-bucket
위 예제 2의 저장 단계에서는 컴파일된 객체 my_compiled_train.pickle에 영향을 미치기 때문에 컴파일러 플래그 LIBTPU_INIT_ARGS 및 learning_rate 내보내기를 포함했습니다 my_compiled_train.pickle. 모델의 크기(예: global_parameter_scale , max_sequence_length 및 per_device_batch )는 compile_train.py 통해 처음 컴파일할 때 고정되어 있으며, 컴파일한 것과 다른 크기로 저장된 컴파일된 개체를 실행하려고 하면 크기 오류가 표시됩니다. 그러나 미묘한 점은 compile_train 실행할 때 학습 속도 일정 도 고정된다는 것입니다. 이는 steps 와 learning_rate 에 의해 결정됩니다. adam_b1 과 같은 최적화 매개변수는 모양이 지정된 객체로만 컴파일러에 전달됩니다. 따라서 실제 값은 컴파일 도중이 아니라 train.py 실행할 때 결정됩니다. 다른 모양(예: per_device_batch )을 전달하는 경우 컴파일된 서명의 예상 모양이 입력된 것과 다르다는 명확한 오류 메시지를 받게 됩니다. compile_topology 통해 요청된 컴파일 대상과 다른 하드웨어에서 실행하려고 하면 컴파일된 장치를 실제 장치에 매핑하는 데 실패했다는 오류가 발생합니다. 컴파일된 것과 다른 XLA 플래그 또는 LIBTPU를 사용하면 컴파일한 환경에서 오류 없이 자동으로 실행될 것입니다. 그러나 이 경우에는 동작이 보장되지 않습니다. 컴파일한 것과 동일한 환경에서 실행해야 합니다.
TPU와는 몇 가지 차이점이 있지만 GPU에서도 사전 컴파일이 지원됩니다.
GPU는 하드웨어 간 컴파일을 지원하지 않습니다. AoT 컴파일을 실행하려면 GPU 호스트가 여전히 필요하지만 단일 GPU 호스트는 동일한 하드웨어의 대규모 클러스터에 대한 프로그램을 컴파일할 수 있습니다.
A3 Cloud GPU의 경우 최대 "슬라이스" 크기는 단일 호스트이고 compile_topology_num_slices 매개변수는 사전 컴파일할 A3 머신 수를 나타냅니다.
이 예에서는 4개의 A3 호스트로 구성된 클러스터를 대상으로 하는 멀티호스트 GPU 컴파일에 사용할 플래그를 보여줍니다.
1단계: AOT 실행 및 컴파일된 함수 저장
# Run the below on a single A3 machine
export XLA_FLAGS="--xla_gpu_enable_async_collectives=true"
python3 MaxText/train_compile.py MaxText/configs/base.yml compile_topology=a3
compile_topology_num_slices=4
compiled_trainstep_file=my_compiled_train.pickle global_parameter_scale=16
attention=dot_product per_device_batch_size=4 steps=10000 learning_rate=1e-3
2단계: train.py를 실행하고 컴파일된 함수 로드
컴파일된 train_step을 로드하려면 compiled_trainstep_file=my_compiled_train.pickle train.py 에 전달하면 됩니다.
# Run the below on each of the 4 target A3 hosts.
export XLA_FLAGS="--xla_gpu_enable_async_collectives=true"
python3 MaxText/train.py MaxText/configs/base.yml run_name=example_load_compile
compiled_trainstep_file=my_compiled_train.pickle
attention=dot_product global_parameter_scale=16 per_device_batch_size=4 steps=10000 learning_rate=1e-3
base_output_directory=gs://my-output-bucket dataset_path=gs://my-dataset-bucket
TPU의 경우와 마찬가지로 컴파일 환경은 실행 환경과 일치해야 하며 이 경우에는 동일한 XLA_FLAGS 설정해야 합니다.
MaxText는 디렉터리에 수집된 로그를 Vertex AI의 Tensorboard 인스턴스에 자동 업로드하는 기능을 지원합니다. 자세한 내용은 사용자 가이드를 따르십시오.
