import torchfrom transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig model_id = "petra312/GLM-4.6-Derestricted-v3-NF4"config = AutoConfig.from_pretrained(model_id, trust_remote_code=True) device_map = { "model.embed_tokens": 0, "model.norm": 0, "model.rotary_emb": 0, "lm_head": 0,}for i in range(config.num_hidden_layers): device_map[f"model.layers.{i}.self_attn"] = 0 device_map[f"model.layers.{i}.input_layernorm"] = 0 device_map[f"model.layers.{i}.post_attention_layernorm"] = 0 if i < config.first_k_dense_replace: device_map[f"model.layers.{i}.mlp"] = 0 else: device_map[f"model.layers.{i}.mlp.gate"] = 0 device_map[f"model.layers.{i}.mlp.shared_experts"] = 0 device_map[f"model.layers.{i}.mlp.experts"] = "cpu" tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)model = AutoModelForCausalLM.from_pretrained( model_id, device_map=device_map, trust_remote_code=True, torch_dtype=torch.bfloat16,)