Birch-san

wsl --update --web-download
wsl --install -d Ubuntu-22.04 --web-download
wsl --setdefault Ubuntu-22.04

sudo apt-get install -y bpftrace bpftrace-dbgsym linux-headers-generic libc6-dev

Birch-san

Some weird memory usage (VRAM) is reported (by torch and by NVML) when using 8-bit AdamW, paged or unpaged.

Here we train llama 2 on 4096-token sequences, using either --optim adamw_8bit or --optim paged_adamw_8bit.
We do a full finetune using qlora.py --full-finetune, with our qlora.py fork, stepwise branch, commit 9a1045d.
We print the memory usage using HF transformers trainer's on_step_end callback. This is after optimizer.step(); model.zero_grad().

One would expect the memory usage at the end of step 1 to be the same as the end of step 2.
Yet for unpaged optimizer: memory usage leaps by 13.2GiB. End of step 1=70.4GiB, end of step 2=81.6GiB.
This appears to be a leap in PyTorch reserved memory only (32.6GiB -> 43.9GiB).

Birch-san

import torch
from transformers import T5ForConditionalGeneration

model: T5ForConditionalGeneration = T5ForConditionalGeneration.from_pretrained('google/t5-v1_1-small')

_inference_mode_context = torch._C._InferenceMode(True)
_inference_mode_context.__enter__()

model.shared.weight.std()
tensor(11.6375)

Birch-san

Running GitHub Copilot VSCode extension against local Code Llama model

Tested on NVIDIA RTX 4090, but these instructions also cover AMD and Mac in case you wanna try those.
This guide assumes you are running Linux (I ran this on Ubuntu).

Before you get excited:

Birch-san

Birch-san

from typing import Optional, NamedTuple
from torch import BoolTensor, arange, meshgrid, clamp
import torch

class Dimensions(NamedTuple):
  height: int
  width: int

def make_neighbourhood_mask(size: Dimensions, size_orig: Dimensions, device='cpu') -> BoolTensor:
  h, w = size

Birch-san

from transformers import (
  AutoConfig,
  AutoTokenizer,
  BitsAndBytesConfig,
  GenerationConfig,
  AutoModelForCausalLM,
  LlamaTokenizerFast,
  PreTrainedModel,
  TextIteratorStreamer,
  StoppingCriteria,

Birch-san

import numpy as np
import math
from numpy.typing import NDArray

# we are trying to make buckets of varying aspect ratios,
# all with about the same area (equivalent to a 512x512 square)
square_side = 512
buckets = 8

widest_aspect: float = math.atan2(1, 2)  # 1/2 = 0.5 aspect ratio

Birch-san

import torch
from typing import Optional
from flash_attn import flash_attn_func, flash_attn_qkvpacked_func
from diffusers.models.attention import Attention

class FlashAttnProcessor:
    r"""
    Processor for implementing memory efficient attention using flash_attn.
    """

Birch-san

import torch
from typing import Optional
from flash_attn import flash_attn_func
from diffusers.models.attention import Attention

class FlashAttnProcessor:
    r"""
    Processor for implementing memory efficient attention using flash_attn.
    """