Update README.md

README.md

@@ -51,25 +51,6 @@ quant_config_4 = BitsAndBytesConfig(
     llm_int8_enable_fp32_cpu_offload=True,
 )
 
-NUM_TRANS_LAYERS = 94
-
-def create_device_map():
-    device_map = {
-        'model.embed_tokens': 0,
-        'model.norm': 0,
-        'lm_head': 0
-    }
-    for start, end, gpu_id in [(0, 4, 0), (4, 22, 1), (22, 40, 2), (40, 58, 3), (58, 76, 4), (76, 94, 5)]:
-        for i in range(start, end):
-            device_map[f'model.layers.{i}'] = gpu_id
-    
-    #for i in range(76, NUM_TRANS_LAYERS):
-    #    device_map[f'model.layers.{i}'] = "cpu"
-
-    return device_map
-
-#device_map = create_device_map()
-
 model = AutoModelForCausalLM.from_pretrained(
     NEW_MODEL_ID, 
     device_map="balanced",