3outeille

# /// script
# requires-python = ">=3.10"
# dependencies = ["huggingface_hub==2.4.0", "matplotlib"]
# ///
"""FSDP vs HSDP across 2 Hugging Face Jobs connected by a network group.

    uv run hsdp_vs_fsdp_hf_jobs.py                      # launch, wait, print the results, write hsdp_vs_fsdp.png
    uv run hsdp_vs_fsdp_hf_jobs.py report <master>,<worker> [<master>,<worker>...]   # plot Jobs that already ran

Launches NNODES Jobs of FLAVOR in one network group. The same Jobs run every layout of LAYOUTS one after the other,

3outeille

Benchmark plot for huggingface/transformers#49813 (dist_moe EP strategy).

3outeille

"""
Reproduce the ~0.5 OOM/year (and actual ~0.6-0.7 OOM/year) training compute trend
using Epoch AI's public frontier models dataset.

Data source: https://epoch.ai/data/ai-models
License: Creative Commons Attribution (CC BY)
"""

import pandas as pd
import numpy as np

3outeille

import torch
import triton
import triton.language as tl

def assert_is_matrix(x):
    if x.ndim != 2:
        raise ValueError(f'Expected 2-tensor but got {x.ndim}-tensor')

def assert_is_vector(x):
    if x.ndim != 1:

3outeille

import os

import torch
import torch.distributed as dist
import lovely_tensors as lt; lt.monkey_patch()


def split_tensor(data: torch.Tensor, dim: int) -> torch.Tensor:
    rank = dist.get_rank()
    world_size = dist.get_world_size()

3outeille

#VERBOSE=0 torchrun --nproc_per_node 3 self_contained_pp_LOC.py
import os, random, numpy as np, torch, torch.nn as nn, torch.distributed as dist, torch.nn.functional as F
from torch.optim import AdamW
from torch.utils.data import DataLoader, DistributedSampler
from datasets import load_dataset
from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer

STEP, local_rank, world_size, verbose = 0, int(os.environ["LOCAL_RANK"]), int(os.environ["WORLD_SIZE"]), os.environ.get("VERBOSE", "0") == "1"

def set_all_seed(seed):

3outeille

3outeille

3outeille

import torch
from torch.nn import functional as F
from torch import distributed as dist
import os
import numpy as np
import random

def set_random_seed(seed: int):
    torch.manual_seed(seed)
    if torch.cuda.is_available():

3outeille

from copy import deepcopy
import torch
from datasets import load_dataset
from torch.optim import SGD
from torch.utils.data import DataLoader
from transformers import AutoModelForCausalLM, AutoTokenizer
import random
import os
import numpy as np