Hugging Face Accelerate

from accelerate import Accelerator

# wandb でログするよう Accelerator オブジェクトに指定します
accelerator = Accelerator(log_with="wandb")

# wandb のパラメーターと必要な設定情報を渡して、wandb run を初期化します
accelerator.init_trackers(
    project_name="my_project", 
    config={"dropout": 0.1, "learning_rate": 1e-2},
    init_kwargs={"wandb": {"entity": "my-wandb-team"}}
    )

...

# accelerator.log() を呼び出して wandb にログします。step は省略可能です
accelerator.log({"train_loss": 1.12, "valid_loss": 0.8}, step=global_step)


# wandb トラッカーが正しく終了するようにします
accelerator.end_training()

wandb_tracker = accelerator.get_tracker("wandb")

wandb_tracker.log_artifact(some_artifact_to_log)