Refactor model loading to support dynamic trust_remote_code and attention implementation options

.gitignore

@@ -7,3 +7,7 @@ atlas/
 *.sqlite
 *.tmp
 .DS_Store
+comparison_post.md
+.gitignore
+model_card_3b.md
+model_card_1.5b.md

analyze_ov_circuits.py

@@ -23,14 +23,14 @@ import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
 
-def _load_model(model_id: str, token: str | None):
-    tokenizer = AutoTokenizer.from_pretrained(model_id, token=token, trust_remote_code=True)
+def _load_model(model_id: str, token: str | None, trust_remote_code: bool = False):
+    tokenizer = AutoTokenizer.from_pretrained(model_id, token=token, trust_remote_code=trust_remote_code)
     model = AutoModelForCausalLM.from_pretrained(
         model_id,
         token=token,
         torch_dtype=torch.bfloat16,
         device_map="cpu",
-        trust_remote_code=True,
+        trust_remote_code=trust_remote_code,
     )
     model.eval()
     return model, tokenizer
@@ -82,9 +82,9 @@ def analyze_head(w_v: np.ndarray, w_o_head: np.ndarray) -> dict:
     }
 
 
-def run(model_id: str, output: Path, token: str | None = None) -> None:
+def run(model_id: str, output: Path, token: str | None = None, trust_remote_code: bool = False) -> None:
     print(f"[ov] loading {model_id}")
-    model, _ = _load_model(model_id, token)
+    model, _ = _load_model(model_id, token, trust_remote_code)
 
     layers = _resolve_layers(model)
     n_layers = len(layers)
@@ -165,9 +165,10 @@ def main() -> None:
     p.add_argument("--model", default=MODEL_ID, help="HF model id")
     p.add_argument("--output", type=Path, default=Path("outputs/ov_circuit_scores.json"))
     p.add_argument("--hf-token", default=os.environ.get("HF_TOKEN"))
+    p.add_argument("--trust-remote", action="store_true", help="Enable trust_remote_code for custom model architectures")
     args = p.parse_args()
 
-    run(args.model, args.output, token=args.hf_token)
+    run(args.model, args.output, token=args.hf_token, trust_remote_code=args.trust_remote)
 
 
 if __name__ == "__main__":

app.py

@@ -171,7 +171,8 @@ def _run_compliance(
             dtype="bfloat16",
             device_map="",
             max_length=128,
-            trust_remote_code=True,
+            trust_remote_code=args.trust_remote,
+            attn_implementation=args.attn_implementation,
         ),
         positive_corpus=CorpusSpec(path=positive, prompt_key=pos_key),
         negative_corpus=CorpusSpec(path=negative, prompt_key=neg_key),
@@ -276,6 +277,8 @@ def _run_sub_zero(
     output: Path,
     report: Path,
     token: str | None,
+    trust_remote: bool = False,
+    attn_implementation: str | None = None,
 ) -> None:
     cmd = [
         sys.executable,
@@ -292,6 +295,10 @@ def _run_sub_zero(
     ]
     if token:
         cmd += ["--hf-token", token]
+    if trust_remote:
+        cmd.append("--trust-remote")
+    if attn_implementation:
+        cmd += ["--attn-implementation", attn_implementation]
     print(f"[sub_zero] running run_sub_zero.py -> {report}")
     subprocess.run(cmd, check=True)
 
@@ -302,7 +309,7 @@ def _atlas_index(atlas_dir: Path) -> None:
     subprocess.run(cmd, check=True)
 
 
-def _run_ov_circuits(model_id: str, output: Path, token: str | None) -> None:
+def _run_ov_circuits(model_id: str, output: Path, token: str | None, trust_remote: bool = False) -> None:
     cmd = [
         sys.executable,
         "analyze_ov_circuits.py",
@@ -311,6 +318,10 @@ def _run_ov_circuits(model_id: str, output: Path, token: str | None) -> None:
         "--output",
         str(output),
     ]
+    if token:
+        cmd += ["--hf-token", token]
+    if trust_remote:
+        cmd.append("--trust-remote")
     print(f"[ov] running analyze_ov_circuits.py -> {output}")
     subprocess.run(cmd, check=True)
 
@@ -329,6 +340,8 @@ def main() -> None:
     p.add_argument("--max-length", type=int, default=128)
     p.add_argument("--batch-size", type=int, default=1)
     p.add_argument("--dtype", default="bfloat16", choices=["bfloat16", "float32"])
+    p.add_argument("--attn-implementation", default=None, choices=["eager", "sdpa", "flash_attention_2"])
+    p.add_argument("--trust-remote", action="store_true", help="Enable trust_remote_code for custom model architectures")
     p.add_argument("--layers", default="all", help="layer spec like '0-15' or 'all'")
     p.add_argument("--sub-zero", action="store_true",
                    help="run the Sub-Zero surgery probe (DAS rotational analysis) and fold it into the atlas")
@@ -386,7 +399,8 @@ def main() -> None:
                 dtype=args.dtype,
                 device_map="",
                 max_length=args.max_length,
-                trust_remote_code=True,
+                trust_remote_code=args.trust_remote,
+                attn_implementation=args.attn_implementation,
             ),
             corpus=CorpusSpec(
                 path=args.corpus,
@@ -414,7 +428,7 @@ def main() -> None:
     # 4. OV-circuit spectral analysis
     print("\n[4/6] OV-circuit analysis")
     ov_report = args.outdir / "ov_circuit_scores.json"
-    _run_ov_circuits(model_id, ov_report, token)
+    _run_ov_circuits(model_id, ov_report, token, trust_remote=args.trust_remote)
 
     # 5. Optional compliance-behaviour extraction
     compliance_report = args.outdir / "compliance_behaviour_scores.json"
@@ -444,6 +458,8 @@ def main() -> None:
             args.outdir / "sub_zero_brain_atlas.json",
             subzero_report,
             token,
+            trust_remote=args.trust_remote,
+            attn_implementation=args.attn_implementation,
         )
     else:
         print("\n[5b/6] skipping Sub-Zero probe (pass --sub-zero to run)")

qwip_atlas/cli.py

@@ -13,10 +13,11 @@ def _extract_local(args: argparse.Namespace) -> None:
         model=ModelSpec(
             model_id=args.model,
             revision=args.revision,
-            trust_remote_code=not args.no_trust_remote_code,
+            trust_remote_code=args.trust_remote,
             dtype=args.dtype,
             device_map=args.device_map,
             max_length=args.max_length,
+            attn_implementation=args.attn_implementation,
         ),
         corpus=CorpusSpec(
             path=Path(args.corpus),
@@ -40,10 +41,11 @@ def _model_spec_from_args(args: argparse.Namespace) -> ModelSpec:
     return ModelSpec(
         model_id=args.model,
         revision=args.revision,
-        trust_remote_code=not args.no_trust_remote_code,
+        trust_remote_code=args.trust_remote,
         dtype=args.dtype,
         device_map=args.device_map,
         max_length=args.max_length,
+        attn_implementation=args.attn_implementation,
     )
 
 
@@ -95,12 +97,13 @@ def main() -> None:
     extract.add_argument("--max-length", type=int, default=512)
     extract.add_argument("--dtype", default="bfloat16", choices=["bfloat16", "bf16", "float16", "fp16", "float32", "fp32"])
     extract.add_argument("--device-map", default="auto", help="Transformers device_map; pass '' to disable")
+    extract.add_argument("--attn-implementation", default=None, choices=["eager", "sdpa", "flash_attention_2"], help="Transformers attention backend")
+    extract.add_argument("--trust-remote", action="store_true", help="Enable trust_remote_code for the model")
     extract.add_argument("--components", default=None, help="Comma-separated subset: mlp,gate,up,attn,heads,q,k,v")
     extract.add_argument("--prompt-key", default="prompt")
     extract.add_argument("--category-key", default="category")
     extract.add_argument("--bucket-key", default="bucket")
     extract.add_argument("--hf-token", default=None)
-    extract.add_argument("--no-trust-remote-code", action="store_true")
     extract.add_argument("--no-truncate", action="store_true")
     extract.set_defaults(func=_extract_local)
 
@@ -119,6 +122,8 @@ def main() -> None:
     compliance_behaviour.add_argument("--max-length", type=int, default=512)
     compliance_behaviour.add_argument("--dtype", default="bfloat16", choices=["bfloat16", "bf16", "float16", "fp16", "float32", "fp32"])
     compliance_behaviour.add_argument("--device-map", default="auto", help="Transformers device_map; pass '' to disable")
+    compliance_behaviour.add_argument("--attn-implementation", default=None, choices=["eager", "sdpa", "flash_attention_2"])
+    compliance_behaviour.add_argument("--trust-remote", action="store_true")
     compliance_behaviour.add_argument("--components", default=None, help="Comma-separated subset: mlp,gate,up,attn,heads,q,k,v")
     compliance_behaviour.add_argument("--prompt-key", default="prompt")
     compliance_behaviour.add_argument("--positive-prompt-key", default=None)
@@ -128,7 +133,6 @@ def main() -> None:
     compliance_behaviour.add_argument("--positive-label", default="positive")
     compliance_behaviour.add_argument("--negative-label", default="negative")
     compliance_behaviour.add_argument("--hf-token", default=None)
-    compliance_behaviour.add_argument("--no-trust-remote-code", action="store_true")
     compliance_behaviour.add_argument("--no-truncate", action="store_true")
     compliance_behaviour.set_defaults(func=_compliance_behaviour_local)
 

qwip_atlas/config.py

@@ -10,10 +10,11 @@ class ModelSpec:
 
     model_id: str
     revision: str | None = None
-    trust_remote_code: bool = True
+    trust_remote_code: bool = False
     dtype: str = "bfloat16"
     device_map: str | None = "auto"
     max_length: int = 512
+    attn_implementation: str | None = None
 
 
 @dataclass(frozen=True)

qwip_atlas/extractors/local_census.py

@@ -55,6 +55,8 @@ def _load_model_and_tokenizer(cfg: AtlasRunConfig, hf_token: str | None):
     }
     if model_spec.device_map:
         kwargs["device_map"] = model_spec.device_map
+    if model_spec.attn_implementation:
+        kwargs["attn_implementation"] = model_spec.attn_implementation
 
     model = AutoModelForCausalLM.from_pretrained(model_spec.model_id, **kwargs)
     model.eval()

qwip_atlas/run_sub_zero.py

@@ -42,23 +42,31 @@ def _torch_dtype(name: str):
     return {"bfloat16": torch.bfloat16, "float32": torch.float32}[name]
 
 
-def _load_model_and_tokenizer(model_id: str, dtype: str, hf_token: str | None):
+def _load_model_and_tokenizer(
+    model_id: str,
+    dtype: str,
+    hf_token: str | None,
+    trust_remote_code: bool = False,
+    attn_implementation: str | None = None,
+):
     import torch
     from transformers import AutoModelForCausalLM, AutoTokenizer
 
     tokenizer = AutoTokenizer.from_pretrained(
-        model_id, trust_remote_code=True, token=hf_token
+        model_id, trust_remote_code=trust_remote_code, token=hf_token
     )
     if tokenizer.pad_token is None:
         tokenizer.pad_token = tokenizer.eos_token
     tokenizer.padding_side = "left"
 
-    model = AutoModelForCausalLM.from_pretrained(
-        model_id,
-        trust_remote_code=True,
-        token=hf_token,
-        torch_dtype=_torch_dtype(dtype),
-    )
+    kwargs = {
+        "trust_remote_code": trust_remote_code,
+        "token": hf_token,
+        "torch_dtype": _torch_dtype(dtype),
+    }
+    if attn_implementation:
+        kwargs["attn_implementation"] = attn_implementation
+    model = AutoModelForCausalLM.from_pretrained(model_id, **kwargs)
     model.eval()
     if torch.cuda.is_available():
         model = model.to("cuda")
@@ -85,6 +93,8 @@ def main() -> None:
     p.add_argument("--layer-limit", type=int, default=None,
                    help="probe only the first N layers (debug/smoke)")
     p.add_argument("--dtype", default="bfloat16", choices=["bfloat16", "float32"])
+    p.add_argument("--attn-implementation", default=None, choices=["eager", "sdpa", "flash_attention_2"])
+    p.add_argument("--trust-remote", action="store_true")
     p.add_argument("--top-k-svs", type=int, default=30,
                    help="bouncer SVs per projection to keep in the report")
     p.add_argument("--all-layers", action="store_true",
@@ -164,7 +174,13 @@ def main() -> None:
     )
 
     print(f"[run_sub_zero] loading {args.model} ({args.dtype}) ...")
-    model, tokenizer = _load_model_and_tokenizer(args.model, args.dtype, args.hf_token)
+    model, tokenizer = _load_model_and_tokenizer(
+        args.model,
+        args.dtype,
+        args.hf_token,
+        trust_remote_code=args.trust_remote,
+        attn_implementation=args.attn_implementation,
+    )
 
     atlas = build_brain_atlas(model, tokenizer, config, cache_path=str(args.output))
     print(f"[run_sub_zero] native brain atlas -> {args.output}")