add poetry env

2026-09-09 11:24:31 +08:00 · 2023-11-12 13:35:39 +08:00
parent ac6be401fe
commit 1800cdf1d0
3 changed files with 3012 additions and 0 deletions
@@ -0,0 +1,41 @@
+[tool.poetry]
+name = "src"
+version = "0.1.0"
+description = ""
+authors = ["wassname <git@wassname.org>"]
+readme = "README.md"
+
+[tool.poetry.dependencies]
+python = ">=3.9,<3.13"
+torch = {version = "^2.1.0+cu118", source = "pytorch"}
+matplotlib = "^3.8.0"
+loguru = "^0.7.2"
+einops = "^0.3.1"
+torchinfo = "^1.8.0"
+accelerate = "^0.24.1"
+peft = "^0.5.0"
+bitsandbytes = {url = "https://github.com/TimDettmers/bitsandbytes/releases/download/0.41.0/bitsandbytes-0.41.0-py3-none-any.whl"}
+transformers = "4.34.0"
+tqdm = "^4.66.1"
+wandb = "^0.12.6"
+ale-py = "^0.8.1"
+pygame = "^2.5.2"
+psutil = "^5.9.6"
+protobuf = "^3.10.0"
+opencv-python = "^4.8.1.78"
+hydra-core = "^1.3.2"
+gym = {extras = ["accept-rom-license"], version = "^0.26.2"}
+torchvision = "^0.16.0"
+
+[[tool.poetry.source]]
+name = "pytorch"
+url = "https://download.pytorch.org/whl/cu118"
+priority = "explicit"
+
+[tool.poetry.group.dev.dependencies]
+ipykernel = "^6.25.2"
+ruff = "^0.1.3"
+
+[build-system]
+requires = ["poetry-core"]
+build-backend = "poetry.core.masonry.api"
@@ -0,0 +1,65 @@
+# 2023-11-12 13:17:35
+
+Try IRIs but with pretrained transformer with LoRA adapter
+
+- [ ] first can I run it
+- [ ] then can I add 3B with adapter...
+
+```sh
+poetry install
+. ./.venv/bin/activate
+python src/main.py env.train.id=BreakoutNoFrameskip-v4 common.device=cuda:0 wandb.mode=online
+```
+
+
+```sh
+# TODO use this code to load a transformer, and other code from my bigvae repo https://github.com/wassname/bigvae_wm
+def load_model(config, device='cuda'):
+    tokenizer = AutoTokenizer.from_pretrained(config.model_name, trust_remote_code=True)
+    tokenizer.padding_side = "left"
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+    bnb_config = BitsAndBytesConfig(
+        load_in_4bit=True,
+        bnb_4bit_compute_dtype=torch.bfloat16,
+        bnb_4bit_quant_type="nf4",
+        bnb_4bit_use_double_quant=True,
+    )
+    base_model = AutoModelForCausalLM.from_pretrained(
+        config.model_name,
+        device_map={"": device},
+        quantization_config=bnb_config,
+        torch_dtype=torch.bfloat16, 
+        trust_remote_code=True
+    )
+    peft_config = peft.LoraConfig(
+        peft.TaskType.CAUSAL_LM,
+        inference_mode=False,
+        r=config.rank,
+        lora_alpha=8,
+        lora_dropout=config.dropout,
+        target_modules=[
+            "self_attn.q_proj",
+            "self_attn.k_proj",
+            "self_attn.v_proj",
+            "self_attn.o_proj",
+            "mlp.gate_proj",
+            "mlp.up_proj",
+            "mlp.down_proj",
+        ],
+    )
+    base_model_peft = peft.get_peft_model(base_model, peft_config)
+    vae_model = BigVAE(
+        base_model_peft, device, peft_config, z_dim=config.z_dim,
+    )
+    if config.start_from:
+        vae_model.load_pretrained(config.start_from)
+    base_model_peft.requires_grad_(False)
+    vae_model.vae_head.requires_grad_(False)
+    vae_model.vae_head.w_d.requires_grad_()
+    router = BigVAERouter(base_model_peft, vae_model, device, peft_config)
+    if config.start_from:
+        router.load_pretrained(config.start_from, is_trainable=True)
+    print(router.model.print_trainable_parameters())
+    router.model.set_adapter("router")
+```