kaivoss/AliceAI-T5-35B-A0.6B-chat-lora-v2

AliceAI-T5-35B-A0.6B · chat + tool-call LoRA, v2

Status: v2, improved but not solved. On the behaviour probes v2 passes 61% (the base model: 11%); on held-out real agent turns it picks the right kind of reply (tool call, text or nothing) 90% of the time and never wrote another turn. It still fails the scenario that motivated it: under a coding-agent prompt, 4 of 10 greetings still get an invented task and a tool call, and after a tool result the model repeats the call instead of answering (0 of 2). It also over-thinks simple questions. Treat it as a research checkpoint, not an agent you can leave unattended.

A LoRA adapter (rank 16, alpha 16) for yandex/AliceAI-T5-35B-A0.6B, a 35B-total / 0.6B-active sparse-MoE encoder-decoder base model with no chat format. It teaches the model to chat, to call tools when a task needs them, to answer in text when it does not, and to say nothing when there is nothing to say.

Adapter	kaivoss/AliceAI-T5-35B-A0.6B-chat-lora-v2 (this repo)
Merged bf16 weights	kaivoss/AliceAI-T5-35B-A0.6B-chat-v2
Fine-tuning example	finetune_chat_lora.py in this repo
v1 (superseded, known issues)	kaivoss/AliceAI-T5-35B-A0.6B-chat-lora

What v1 got wrong, and what changed

v1 ran inside an agent harness and misbehaved: a bare "hi" produced 17 to 22 parallel tool calls, "don't read any skills" was ignored, replies were sometimes empty or endless. Three causes were found and fixed here:

1. Untrained delimiter tokens. v1's format used the vocabulary's [INST_START], [REPLY_END], [SYSTEM_*], [TOOL_RESULT_*] tokens. Their embeddings in the base model are untrained placeholders, identical to the unused <SPAN#5000>. The model could not tell roles apart or emit [REPLY_END], and LoRA on attention cannot teach embeddings. v2 uses only tokens the base model trained: plain-text role markers and </s> as the end of a turn. placeholder_ids.json lists all 5,832 placeholder tokens, and a unit test fails if the chat format ever uses one.
2. Lopsided data. Outside the Tool-Use subsplit about 98% of v1's training turns were tool calls. v2 mixes in plain chat, tools-present-but-not-needed cases, "don't use tools" instructions and harness-style prompts (table below).
3. Greedy decoding. The serving shim defaulted to temperature 0; the harness sends none. The shim now samples, merges duplicate calls, caps parallel calls and never lets a turn continue as the user.

Chat template

The base repo ships no chat template; chat_template.jinja defines one inside the model's UL2 frame:

encoder:  [_S_]<|system|>…\n<|user|>…\n<|assistant|>…\n<|tool|>…\n<|assistant|><SPAN#0>
decoder:  <s><SPAN#0>  →  <think>reasoning</think>text\n<tool_call>name<arg_key>k</arg_key><arg_value>v</arg_value></tool_call></s>

Tools go in the system message as text in the GLM format (# Tools … <tools>{json per line}</tools> …), produced by glm_tools.py. Reasoning is rendered only for the turn being predicted. An empty reply is just </s>. Pass alice_frame=False to apply_chat_template when a runtime adds [_S_] and <SPAN#0> itself (vLLM plugin below).

Use

import torch
from huggingface_hub import hf_hub_download
from peft import PeftModel
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer

for f in ("alice_chat.py",):
    hf_hub_download("kaivoss/AliceAI-T5-35B-A0.6B-chat-lora-v2", f, local_dir=".")
from alice_chat import clean_decode, parse_output

BASE, ADAPTER = "yandex/AliceAI-T5-35B-A0.6B", "kaivoss/AliceAI-T5-35B-A0.6B-chat-lora-v2"
tok = AutoTokenizer.from_pretrained(ADAPTER, trust_remote_code=True)          # carries the chat template
model = AutoModelForSeq2SeqLM.from_pretrained(BASE, trust_remote_code=True, dtype=torch.bfloat16,
                                              attn_implementation="eager", device_map={"": "cuda:0"})
model = PeftModel.from_pretrained(model, ADAPTER).eval()     # or load the merged repo directly

messages = [{"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": "Give me one fact about Zürich."}]
prompt = tok.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
src = tok(prompt, add_special_tokens=False, return_tensors="pt").input_ids.to(model.device)
span = tok.convert_tokens_to_ids("<SPAN#0>")
dec = torch.tensor([[model.config.decoder.bos_token_id, span]], device=model.device)
out = model.generate(input_ids=src, attention_mask=torch.ones_like(src), decoder_input_ids=dec,
                     do_sample=True, temperature=0.6, top_p=0.95, max_new_tokens=512)
gen = out[0, 2:].tolist()
gen = gen[: gen.index(tok.eos_token_id)] if tok.eos_token_id in gen else gen
reasoning, content, tool_calls = parse_output(clean_decode(tok, gen))

Use sampling, not greedy: greedy decoding looped in testing.

Serving with vLLM and an OpenAI-compatible client

The community plugin (AntonProkopyev/AliceAI-T5-35B-A0.6B-vLLM) loads the merged repo unchanged, but it only serves /v1/completions and adds the UL2 frame itself. alice_openai_shim.py sits in front of it and provides /v1/chat/completions with tools, reasoning and streaming (simulated), plus the guardrails above. serve_smoke.py exercises it with the openai Python client.

Training

Method	LoRA (bf16 base; not QLoRA, not full fine-tuning)
Rank / alpha / dropout	16 / 16 / 0.05 (scale 1.0)
Targets	q_proj, k_proj, v_proj, o_proj in encoder self-attn, decoder self-attn and decoder cross-attn (experts are fused parameters, not nn.Linear)
Steps	200 (batch 1 × 8 accumulation, each example seen once)
Optimiser	AdamW, lr 2e-4, cosine, 10 warmup steps, bf16, gradient checkpointing
Context	up to 32,768 tokens; longer examples are skipped per turn, never truncated, and turn choice is fit-aware so no row is lost
Attention	SDPA (the remote code does not declare it; enable_sdpa() sets the flag, checked against eager on padded batches)
Hardware	1 × H200

Data

One example per assistant turn: the encoder reads everything before it, the decoder learns that turn, then </s>.

category	examples	mean tokens	max tokens	reply kinds
agent_action	920	16,339	31,325	{'action': 920}
agent_plain	280	10,959	30,087	{'plain': 280}
chat_plain	300	512	2,902	{'plain': 300}
chat_tools	240	4,805	8,548	{'plain': 240}
eval_agent	100	15,040	30,922	{'action': 87, 'plain': 13}
eval_chat	40	938	3,385	{'plain': 40}
harness	100	2,181	4,476	{'plain': 100}
silent	160	13,413	29,603	{'silent': 160}

Sources: real agent trajectories (UltraData-SFT-Agent-2609, GLM-formatted copy linked above), plain Q&A from UltraData-SFT-2605 (English IF and Knowledge, with and without reasoning), greetings and system-prompted chats from smoltalk, and a sanitized capture of the pi coding-agent prompt and tool schemas (harness_pi.json). Targets are filtered to at most 8 parallel calls and 2,000 reasoning tokens. Held-out rows are never trained on.

Evaluation

Behaviour probes (probes.py) come from the failures seen with v1: greetings and "don't use tools" under a pi-style prompt with and without a long skills list, plain questions with tools present, real actions, answering after a tool result, and silence. A probe passes only if every expectation holds, the reply finished, and it never wrote another turn.

probe category	base greedy	base sampled	LoRA greedy	LoRA sampled
action	0/3	0/3	2/3	2/3
after_tool	0/2	0/2	0/2	0/2
greeting	0/10	0/10	6/10	6/10
negation	0/4	0/4	2/4	2/4
plain_chat	3/4	3/4	4/4	4/4
plain_question	0/3	0/3	2/3	2/3
silent	0/2	0/2	1/2	1/2
all probes passed	11%	11%	61%	61%

Held-out real turns (adapter only):

n = 60 held-out rows, greedy, 1,024 new tokens.

metric	adapter
generations cut off at the token cap (not judged below)	18.3%
right kind of reply (tool call / text / nothing)	89.8%
... when the right reply was a tool call	97.1%
... when the right reply was text	73.3%
... when the right reply was nothing	n/a
first tool name matches the real one	45.5%
wrote another turn (`<	user
turns with duplicate tool calls	1.7%
most tool calls in one turn	11

These are small samples (a few dozen probes, n held-out rows) and greedy or single-sample generations. They are a check for the specific failures above, not a capability benchmark. Tool-call arguments are not checked, only the tool name.

Limitations

• Small budget. 200 steps is 1,600 examples: enough to move behaviour, not to add agentic competence the base lacks.
• Attention-only adaptation. Experts and embeddings are untouched.
• Tokenizer quirk. The tokenizer inserts a space after every newline when decoding; use clean_decode from alice_chat.py (important for code in tool arguments).
• Data. Reasoning traces come from whatever model produced the upstream data, not from AliceAI. Plain chat comes from English-only sources. No safety tuning was done: the model can and will run bash commands if a harness lets it.
• Not verified under vLLM. The merged weights and framing were verified for v1; v2's serving path has the same structure but see the repository history for what was actually re-run.

Licence and attribution

Apache-2.0, matching the base model (© Yandex LLC) and the data (UltraData: arXiv:2602.09003; smoltalk: Apache-2.0 for its new subsets).