Usar o diagnóstico de ML com o MaxText

O MaxText é uma biblioteca de modelo de linguagem grande (LLM) de código aberto criada para ser executada em TPUs e GPUs para treinamento de modelos. O SDK ML Diagnostics vem pré-integrado ao MaxText. É possível ativar o ML Diagnostics com o MaxText usando a flag managed-mldiagnostics.

Quando essa opção estiver ativada, você poderá fazer o seguinte:

• Crie uma execução gerenciada do MachineLearning com todas as configurações do MaxText.
• Faça upload dos rastreamentos de criação de perfil se a criação de perfil estiver ativada pela flag profiler="xplane".
• Faça upload das métricas de treinamento em um intervalo definido usando a flag log_period.

Use as flags a seguir ao usar o ML Diagnostics com o MaxText:

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

Para ativar o ML Diagnostics no MaxText, mude o arquivo de configuração da sua execução ou transmita as flags da linha de comando. Por exemplo, transmita as seguintes flags com o comando 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"

Neste exemplo, definir a flag upload_all_profiler_results=True captura perfis de vários hosts de todos os hosts.