diff --git a/laya/agent.py b/laya/agent.py index b0fb78a..0d88c86 100644 --- a/laya/agent.py +++ b/laya/agent.py @@ -151,19 +151,20 @@ def _load_tokenizer(tok_dir: str, cfg: Dict) -> Any: """ from transformers import AutoTokenizer + remote = bool(cfg.get("trust_remote_code", False)) if not os.path.isdir(tok_dir): - return AutoTokenizer.from_pretrained(cfg.get("encoder")) + return AutoTokenizer.from_pretrained(cfg.get("encoder"), trust_remote_code=remote) try: stamp = (os.path.abspath(tok_dir), os.path.getmtime(os.path.join(tok_dir, "tokenizer_config.json"))) except OSError: - return AutoTokenizer.from_pretrained(tok_dir) + return AutoTokenizer.from_pretrained(tok_dir, trust_remote_code=remote) with _TOKENIZERS_LOCK: cached = _TOKENIZERS.get(stamp) if cached is not None: return cached - tokenizer = AutoTokenizer.from_pretrained(tok_dir) + tokenizer = AutoTokenizer.from_pretrained(tok_dir, trust_remote_code=remote) _TOKENIZERS[stamp] = tokenizer return tokenizer diff --git a/laya/common.py b/laya/common.py index 6c4c426..2be1aec 100644 --- a/laya/common.py +++ b/laya/common.py @@ -201,12 +201,26 @@ def _apply_rope_config(ecfg) -> None: def build_model(cfg: Dict, encoder_dir: Optional[str] = None, pretrained: bool = True) -> DecisionModel: from transformers import AutoConfig, AutoModel + # Encoders whose modeling code lives on the Hub (e.g. NeoBERT) must opt in via the config. + remote = bool(cfg.get("trust_remote_code", False)) if not pretrained or (encoder_dir and os.path.exists(encoder_dir)): - ecfg = AutoConfig.from_pretrained(encoder_dir or cfg["encoder"]) + ecfg = AutoConfig.from_pretrained(encoder_dir or cfg["encoder"], trust_remote_code=remote) _apply_rope_config(ecfg) - enc = AutoModel.from_config(ecfg, attn_implementation="sdpa") + enc = AutoModel.from_config(ecfg, attn_implementation="sdpa", trust_remote_code=remote) else: - enc = AutoModel.from_pretrained(cfg["encoder"], attn_implementation="sdpa") + enc = AutoModel.from_pretrained(cfg["encoder"], attn_implementation="sdpa", trust_remote_code=remote) + # NeoBERT computes its RoPE tables as non-persistent buffers in __init__; transformers 5 builds models on + # the meta device, so they come back as uninitialised memory and every forward pass returns NaN. + if hasattr(enc, "freqs_cos"): + import sys + + precompute = sys.modules[type(enc).__module__].precompute_freqs + c = enc.config + enc.freqs_cos, enc.freqs_sin = precompute(c.hidden_size // c.num_attention_heads, c.max_length) + # Some encoder configs pin a half-precision dtype (NeoBERT ships bfloat16), which transformers 5 honours. Keep + # master weights in float32: bf16 weights round most optimizer updates to zero, and CPU inference in bf16 is ~8x slower. + # Mixed precision is applied by autocast at train/inference time instead. + enc = enc.float() return DecisionModel(enc, cfg.get("head_layers", 2), len(cfg.get("act_costs", {})) + 1)