fp16

p16.py → fp16.py

@@ -9,6 +9,7 @@ from tqdm import tqdm
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from transformers.generation.stopping_criteria import StoppingCriteria, StoppingCriteriaList
 
+# Make sure to set your model output directory and make sure it has 755 permissions.
 MODEL_ID = "openai/gpt-oss-120b"
 OUTPUT_DIR = os.environ.get("OUTPUT_DIR", "./fp16/gpt-oss-120b-fp16")
 
@@ -19,11 +20,12 @@ os.environ["TOKENIZERS_PARALLELISM"] = "false"
 tok = AutoTokenizer.from_pretrained(MODEL_ID, use_fast=True)
 
 # 3. load model in fp16
-max_memory = {0: "17GiB", 1: "17GiB", 2: "17GiB", 3: "17GiB", 4: "17GiB", 5: "17GiB", "cpu": "196GiB"}
+# Make sure to change these max_memory settings.
+max_memory = {0: "17GiB", 1: "17GiB", 2: "17GiB", 3: "17GiB", 4: "17GiB", 5: "17GiB", 6: "17GiB", 7: "17GiB", "cpu": "196GiB"}
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_ID,
     torch_dtype=torch.float16,
-    device_map="auto",
+    device_map="sequential",
     low_cpu_mem_usage=True,
     max_memory=max_memory,
     offload_folder="./offload_cache",
@@ -98,7 +100,7 @@ for _m in model.modules():
 from transformers.models.gpt_bigcode import modeling_gpt_bigcode
 modeling_gpt_bigcode.GPTBigCodeModel._check_hidden_states_dtype = lambda *_, **__: None
 
-# 5. quick demo
+# 5. inference to verify functionality
 if __name__ == "__main__":
     prompt = "Explain quantum supremacy in one paragraph."
     inputs = tok(prompt, return_tensors="pt").to(model.device)