MaxText와 함께 ML 진단 사용

MaxText는 모델 학습을 위해 TPU 및 GPU에서 실행되도록 빌드된 오픈소스 대규모 언어 모델 (LLM) 라이브러리입니다. ML Diagnostics SDK는 MaxText와 사전 통합되어 있습니다. managed-mldiagnostics 플래그를 사용하여 MaxText로 ML 진단을 사용 설정할 수 있습니다.

이 옵션을 사용 설정하면 다음 작업을 수행할 수 있습니다.

• 모든 MaxText 구성으로 관리형 MachineLearning 실행을 만듭니다.
• profiler="xplane" 플래그로 프로파일링이 사용 설정된 경우 프로파일링 트레이스를 업로드합니다.
• log_period 플래그를 사용하여 정의된 간격으로 학습 측정항목을 업로드합니다.

MaxText와 함께 ML 진단을 사용할 때는 다음 플래그를 사용하세요.

managed_mldiagnostics: True  # Enable the managed diagnostics
managed_mldiagnostics_run_group: GROUP_NAME  # Used to group multiple runs. (Optional)

MaxText에서 ML 진단을 사용 설정하려면 실행의 구성 파일을 변경하거나 명령줄에서 플래그를 전달하면 됩니다. 예를 들어 MaxText.train 명령어와 함께 다음 플래그를 전달합니다.

python3 -m MaxText.train \
  src/MaxText/configs/base.yml \
  run_name="demo-mldiagnostics-run-2" \
  model_name="<your_chosen_model>" \
  base_output_directory=gs://<your_gcs_folder>/ \
  dataset_type=synthetic \
  steps=100 \
  log_period=10 \
  profiler=xplane \
  upload_all_profiler_results=True \
  managed_mldiagnostics=True \
  managed_mldiagnostics_run_group="demo-mldiagnostics-group"

이 예에서는 upload_all_profiler_results=True 플래그를 설정하면 모든 호스트에서 다중 호스트 프로필이 캡처됩니다.