diff --git a/laya/agent.py b/laya/agent.py index 70163bf..b7b9a0b 100644 --- a/laya/agent.py +++ b/laya/agent.py @@ -186,7 +186,8 @@ class Agent: self.device = torch.device("cpu") tok_dir = os.path.join(model_dir, "tokenizer") - self.tok = AutoTokenizer.from_pretrained(tok_dir if os.path.exists(tok_dir) else self.cfg.get("encoder")) + self.tok = AutoTokenizer.from_pretrained(tok_dir if os.path.exists(tok_dir) else self.cfg.get("encoder"), + trust_remote_code=bool(self.cfg.get("trust_remote_code", False))) enc_dir = os.path.join(model_dir, "encoder") # The checkpoint supplies every parameter; skip random/base-model weights. diff --git a/laya/common.py b/laya/common.py index 950c41d..6a633f7 100644 --- a/laya/common.py +++ b/laya/common.py @@ -139,11 +139,25 @@ class DecisionModel(nn.Module): 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"]) - enc = AutoModel.from_config(ecfg, attn_implementation="sdpa") + ecfg = AutoConfig.from_pretrained(encoder_dir or cfg["encoder"], trust_remote_code=remote) + 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)