NousResearch/atropos
1
0
Fork 0
mirror of https://github.com/NousResearch/atropos.git synced 2026-04-24 17:04:55 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

readme updates

Browse source
This commit is contained in:
Jai Suphavadeeprasit 2026-01-27 17:13:18 -05:00
parent 998c062709
commit 41cab9d52a
2 changed files with 47 additions and 15 deletions
Download patch file Download diff file Expand all files Collapse all files

50
example_trainer/model.py
View file

@ -65,9 +65,18 @@ def load_model_and_tokenizer(
else:
# Legacy mode: load full model
print("[Setup] Loading model for legacy mode...")
model = AutoModelForCausalLM.from_pretrained(
config.model_name, torch_dtype=torch.bfloat16
)
try:
model = AutoModelForCausalLM.from_pretrained(
config.model_name,
torch_dtype=torch.bfloat16,
attn_implementation="flash_attention_2",
)
print("[Setup] Using Flash Attention 2")
except Exception as e:
print(f"[Setup] Flash Attention 2 not available: {e}")
model = AutoModelForCausalLM.from_pretrained(
config.model_name, torch_dtype=torch.bfloat16
)
model.to(config.device)
# Enable gradient checkpointing
@ -121,9 +130,18 @@ def _load_model_with_lora(config: TrainingConfig) -> torch.nn.Module:
raise RuntimeError("PEFT library not available. Install with: pip install peft")
print("[Setup] Loading base model for LoRA mode...")
base_model = AutoModelForCausalLM.from_pretrained(
config.model_name, torch_dtype=torch.bfloat16
)
try:
base_model = AutoModelForCausalLM.from_pretrained(
config.model_name,
torch_dtype=torch.bfloat16,
attn_implementation="flash_attention_2",
)
print("[Setup] Using Flash Attention 2")
except Exception as e:
print(f"[Setup] Flash Attention 2 not available: {e}")
base_model = AutoModelForCausalLM.from_pretrained(
config.model_name, torch_dtype=torch.bfloat16
)
base_model.to(config.device)
# Determine target modules
@ -215,11 +233,23 @@ def _attach_to_vllm_shared_tensors(
model_config = AutoConfig.from_pretrained(config.model_name)
# Create empty model on meta device (no memory allocation)
# Use Flash Attention 2 to match vLLM's attention implementation more closely
# This reduces logprob differences from ~10-15% to ~1-2%
with torch.device("meta"):
model = AutoModelForCausalLM.from_config(
model_config,
torch_dtype=torch.bfloat16,
)
try:
model = AutoModelForCausalLM.from_config(
model_config,
torch_dtype=torch.bfloat16,
attn_implementation="flash_attention_2",
)
print("[Setup] Using Flash Attention 2 (matches vLLM)")
except Exception as e:
print(f"[Setup] Flash Attention 2 not available ({e}), using default attention")
print("[Setup] WARNING: This may cause ~10-15% logprob difference with vLLM")
model = AutoModelForCausalLM.from_config(
model_config,
torch_dtype=torch.bfloat16,
)
param_names = list(model.state_dict().keys())
print(f"[Setup] Model architecture has {len(param_names)} parameters", flush=True)

12
example_trainer/trainers.py
View file

@ -291,17 +291,19 @@ def train_shared_vllm(config: TrainingConfig):
if trainer_logprobs and vllm_logprobs:
for i, (vlp, tlp) in enumerate(zip(vllm_logprobs, trainer_logprobs)):
diff = abs(vlp - tlp)
status = "✓" if diff < 0.1 else "✗"
status = "✓" if diff < 0.25 else "⚠" # 0.25 threshold accounts for impl differences
print(f" Token {i}: vLLM={vlp:.4f}, Trainer={tlp:.4f}, diff={diff:.4f} {status}")
mean_diff = sum(abs(v-t) for v,t in zip(vllm_logprobs, trainer_logprobs)) / len(trainer_logprobs)
print(f" Mean diff: {mean_diff:.4f}")
if mean_diff < 0.1:
print(f" ✓ WEIGHTS ARE SHARED CORRECTLY!")
if mean_diff < 0.05:
print(f" ✓ PERFECT ALIGNMENT - weights shared and same compute path")
elif mean_diff < 0.25:
print(f" ✓ WEIGHTS ARE SHARED (diff {mean_diff:.2f} is due to different forward pass implementations)")
print(f" vLLM uses Flash Attention, trainer uses HuggingFace - small diff is expected!")
else:
print(f" ✗ WEIGHTS ARE NOT SHARED - IPC attachment may have failed!")
print(f" The trainer may have copies of weights, not shared memory.")
print(f" ⚠ Large diff ({mean_diff:.2f}) - may indicate issue with weight sharing")
else:
print(f" vLLM request failed: {response.status_code}")