Picotron Documentation

repository·main·Indexed 25 days ago

https://github.com/huggingface/picotron

A minimalist, educational repository for pre-training Llama-like models using 4D Parallelism, including Data, Tensor, Pipeline, and Context parallel. Designed for researchers and students, it supports GPU and CPU training and provides tools for configuration generation and Slurm job submission.

Tokens
355
Snippets
1
Records
4
Agent score
31%

What's inside Picotron

  1. Quick start: GPU training

    main

    To start training on GPUs, you first generate a configuration JSON file using create_config.py and then execute the training script using torchrun.

    Prerequisites:

    • A Hugging Face token to download models.

    Workflow:

    1. Generate Config: Use create_config.py to define your experiment parameters (model name, parallelism degrees, sequence length, etc.). By default, this creates a JSON file in the specified --out_dir.
    2. Run Training: Use torchrun pointing to the generated config file.

    Parallelism Options:

    • Data Parallelism (DP): Controlled via --dp.
    • 3D Parallelism: Combines Tensor Parallelism (--tp), Pipeline Parallelism (--pp), and Data Parallelism (--dp). You can specify the pipeline engine using --pp_engine (e.g., 1f1b).
    1. Slurm: For cluster environments, use submit_slurm_jobs.py to submit the jobs defined in your config directory.