goliath-120b-jlens-wip

Only trained 14 steps with dim_batch=14 to fit 140gb vram

Samples: https://huggingface.co/gghfexp/goliath-120b-jlens-wip/tree/main/samples

Left the checkpoint in case someone else wants to resume it:

# Configure
# ---------

#Model/Repo setup shit
MODEL_NAME = "alpindale/goliath-120b"
CHECKPOINT="goliath-120b_it_ckpt.pt"
OUTPUT="goliath-120b_it_jlens.pt"
REPO="goliath-120b-it-jacobian-lens"

#Bits
load_in_4bit=True

# ---------


# Load The Model
if load_in_4bit:
  import torch
  import transformers

  quant_config = transformers.BitsAndBytesConfig(
      load_in_4bit=True,
      bnb_4bit_quant_type="nf4",              # apparently NF4 is better than plain fp4
      bnb_4bit_compute_dtype=torch.bfloat16,  # compute in bf16
      bnb_4bit_use_double_quant=True,         # apparently small extra memory saving but didn't do anything for me
  )

  hf_model = transformers.AutoModelForCausalLM.from_pretrained(
      MODEL_NAME,
      quantization_config=quant_config,
      attn_implementation="sdpa",
      device_map="auto",
  )#.cuda()
  tokenizer = transformers.AutoTokenizer.from_pretrained(MODEL_NAME)
else:
  print("16bit")
  import torch
  import transformers

  hf_model = transformers.AutoModelForCausalLM.from_pretrained(
      MODEL_NAME, dtype=torch.bfloat16, attn_implementation="sdpa"
  ).cuda()
  tokenizer = transformers.AutoTokenizer.from_pretrained(MODEL_NAME)

hf_model.gradient_checkpointing_enable()
hf_model.config.use_cache = False
hf_model.requires_grad_(False)
import jlens
jlens.configure_logging()
model = jlens.from_hf(hf_model, tokenizer)
model

# Jlens
from jlens.examples import load_wikitext_prompts


#Gemma-4 modded (fussy model)
#prompts = load_wikitext_prompts(n_prompts=128,tokenizer=tokenizer)


prompts = load_wikitext_prompts(n_prompts=128)

#dim_batch >= 15 == OOM on 140GB
lens = jlens.fit(
    model, prompts, dim_batch=14, max_seq_len=128, checkpoint_path=CHECKPOINT
)

lens.save(OUTPUT)