Guitaricet

class TransformerLayer(nn.Module):
    def __init__(self, dim=496, heads=4, ffn_dim=1984):
        super().__init__()
        self.attention = Attention(dim, heads=heads)

        self.fc = nn.Sequential(
            nn.Linear(dim, ffn_dim),
            nn.ReLU(),  # chose your favorite nonlinearity here
            nn.Linear(ffn_dim, dim),
        )

Guitaricet

class LiteTransformerLayer(nn.Module):
    def __init__(self, dim=496, heads=4, kernel_size=4):
        super().__init__()
        assert dim % 2 == 0

        self.attention = Attention(dim // 2, heads=heads)
        self.cnn = LightweightConv(dim // 2, kernel=kernel_size)  # or Dynamic conv

        self.fc = nn.Sequential(
            nn.Linear(dim, dim),

Guitaricet

Guitaricet

import os
import sys
import logging

logging.basicConfig(
    format="%(asctime)s | %(levelname)s | %(name)s | %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
    level=logging.INFO,
    stream=sys.stdout,
)

Guitaricet

# torchrun --nproc_per_node 2 generate_multigpu.py
# tested on torch==1.12.1 and transformers==4.21
import os
import json

import torch
import torch.distributed as dist

import transformers
import datasets

Guitaricet

import torch


class NamedShape:
    """A convenience class to make beautifully named shapes."""
    def __init__(self, tensor) -> None:
        self.names = tensor.names
        self.shape = tensor.shape

    def __repr__(self) -> str:

Guitaricet

from torch.nn import ModuleList


def _addindent(s_, numSpaces):
    s = s_.split('\n')
    # don't do anything for single-line stuff
    if len(s) == 1:
        return s_
    first = s.pop(0)
    s = [(numSpaces * ' ') + line for line in s]

Guitaricet

import netrc
import os

def is_wandb_logged_in():
    netrc_path = os.path.expanduser("~/.netrc")
    if not os.path.exists(netrc_path):
        return False
    
    auth = netrc.netrc(netrc_path).authenticators("api.wandb.ai")
    return bool(auth)