Goliath 120b Jlens Wip

--- base_model: - alpindale/goliath-120b --- 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: ```python # 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) ```