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

3outeille

import torch
from tqdm import tqdm
from transformers import AutoModelForCausalLM, AutoTokenizer
from datasets import load_dataset

# Model
device = "cpu"

device_map = {

3outeille

• Matmul benchmark of Group-ordering vs Row-major ordering on A100 => No significant improvment over row-major ordering

https://triton-lang.org/master/getting-started/tutorials/03-matrix-multiplication.html#l2-cache-optimizations

matmul-performance:
        M  group_ordering  row_major_ordering
0   256.0        3.640889            3.640889
1   384.0       11.059200           12.288000
2   512.0       23.831273           23.831273
3   640.0       39.384616           39.384616
4 768.0 58.982401 58.982401