ZacharyZhang-NY
9cb46c18d3
The 500-line rule and the one-line-comment rule were enforced only from
memory, so they drifted. bench/filesize_check.py now checks the budget on
every tracked source file, and it passes.
Split, with behaviour proven unchanged by comparing each moved function body
against HEAD byte for byte (53 training bodies, 21 benchmark bodies, 0
differences):
db.py 503 -> 396 store_search.py 105
bench/run.py 609 -> 351 query_set.py 80, report.py 211
benchmarks.py 619 -> 424 benchmark_scoring.py 262
training.py 1630 -> 94 training_{common,encode,metrics,datasets,lora,retrieval}
TrainingService stays as a thin facade for the CLI; the facade's parameter
lists are pinned against the module functions by inspect, which is what
catches a leaked self.
Deleted the eight top-level functions no file called, after fixing the scan
that had reported 201 false positives by skipping the defining file. The
v1/v2/v3 sweep found nothing to delete and says so: every remaining version
tag is an on-disk format marker, a migration step, a fixture name, or a
third-party model name.
Collapsed 39 multi-line comment blocks (86 lines) to one line each, keeping
the measured facts and error strings that justify the code.
269 tests pass; ruff clean; mypy 24 errors in 9 files against a 24-in-6
baseline, net zero; pyright 4, the same four HEAD measures.
359 lines
15 KiB
Python
359 lines
15 KiB
Python
"""Train the embedding adapter: pairs in, cosine ranking out, weights measured after."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import time
|
|
from datetime import UTC, datetime
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from .config import Settings
|
|
from .db import Database
|
|
from .schemas import AdapterManifest, TrainingAdapterKind
|
|
from .training_common import (
|
|
AUTOCAST_DTYPE,
|
|
GRADIENT_CHECKPOINTING_MODE,
|
|
TRAINING_FRAMEWORK,
|
|
TRAINING_METHOD,
|
|
TRAINING_PRECISION,
|
|
TRAINING_PROOF_VERSION,
|
|
adapter_id,
|
|
relative_artifacts,
|
|
sha256_file,
|
|
sha256_tree,
|
|
write_trace,
|
|
write_training_summary,
|
|
)
|
|
from .training_datasets import validate_dataset
|
|
from .training_encode import encode_retrieval_batch, pool_embeddings
|
|
from .training_lora import write_lora_config
|
|
from .training_metrics import (
|
|
activate_unsloth_training,
|
|
bitsandbytes_gaudi_probe_guard,
|
|
cuda_peak_memory_mb,
|
|
enable_gradient_checkpointing,
|
|
gradient_l2_norm,
|
|
infer_lora_targets,
|
|
make_optimizer,
|
|
model_type_for_proof,
|
|
optimizer_family,
|
|
package_version,
|
|
require_cuda_bf16,
|
|
safetensor_weight_stats,
|
|
trainable_l2_norm,
|
|
trainable_named_parameters,
|
|
trainable_state_cpu,
|
|
trainable_state_sha256,
|
|
trainable_update_l2,
|
|
training_contract,
|
|
training_environment,
|
|
)
|
|
|
|
|
|
def train_retrieval_lora(
|
|
settings: Settings,
|
|
database: Database,
|
|
dataset_path: Path,
|
|
output_dir: Path,
|
|
max_steps: int,
|
|
rank: int = 8,
|
|
alpha: int = 16,
|
|
max_length: int | None = None,
|
|
learning_rate: float | None = None,
|
|
cartridge_id: str | None = None,
|
|
) -> AdapterManifest:
|
|
samples = validate_dataset(dataset_path)
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
max_length = max_length or settings.train_max_length
|
|
learning_rate = learning_rate or settings.train_learning_rate
|
|
started_at = datetime.now(UTC)
|
|
wall_start = time.perf_counter()
|
|
trace_path = output_dir / "training_trace.jsonl"
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
activate_unsloth_training()
|
|
from peft import LoraConfig, TaskType, get_peft_model
|
|
from transformers import AutoModel, AutoProcessor
|
|
|
|
model_id = settings.embedding_model_id
|
|
require_cuda_bf16(torch)
|
|
processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
|
|
model = AutoModel.from_pretrained(
|
|
model_id,
|
|
dtype=torch.bfloat16,
|
|
device_map={"": 0},
|
|
trust_remote_code=True,
|
|
)
|
|
if hasattr(model.config, "use_cache"):
|
|
model.config.use_cache = False
|
|
target_modules = infer_lora_targets(model, TrainingAdapterKind.hme_retrieval)
|
|
gradient_checkpointing = enable_gradient_checkpointing(model)
|
|
with bitsandbytes_gaudi_probe_guard():
|
|
model = get_peft_model(
|
|
model,
|
|
LoraConfig(
|
|
r=rank,
|
|
lora_alpha=alpha,
|
|
target_modules=target_modules,
|
|
lora_dropout=0.0,
|
|
bias="none",
|
|
task_type=TaskType.FEATURE_EXTRACTION,
|
|
),
|
|
)
|
|
unsloth_model_type = model_type_for_proof(model)
|
|
trainable_params, total_params = model.get_nb_trainable_parameters()
|
|
optimizer, optimizer_name = make_optimizer(model, learning_rate)
|
|
model.train()
|
|
|
|
losses: list[float] = []
|
|
optimizer_updates: list[dict[str, Any]] = []
|
|
trainable_tensors = trainable_named_parameters(model)
|
|
initial_trainable_sha = trainable_state_sha256(trainable_tensors)
|
|
initial_trainable_l2 = trainable_l2_norm(trainable_tensors)
|
|
batch_size = min(4, len(samples))
|
|
environment = training_environment(torch)
|
|
contract = training_contract(
|
|
gradient_checkpointing=gradient_checkpointing,
|
|
input_require_grads=False,
|
|
optimizer_name=optimizer_name,
|
|
unsloth_model_type=unsloth_model_type,
|
|
)
|
|
with trace_path.open("w", encoding="utf-8", newline="\n") as trace:
|
|
write_trace(
|
|
trace,
|
|
{
|
|
"event": "run_start",
|
|
**contract,
|
|
"started_at": started_at.isoformat(),
|
|
"adapter_kind": TrainingAdapterKind.hme_retrieval.value,
|
|
"base_model": model_id,
|
|
"dataset_path": str(dataset_path.resolve()),
|
|
"dataset_sha256": sha256_file(dataset_path),
|
|
"sample_count": len(samples),
|
|
"max_steps": max_steps,
|
|
"max_length": max_length,
|
|
"rank": rank,
|
|
"alpha": alpha,
|
|
"target_modules": target_modules,
|
|
"trainable_params": int(trainable_params),
|
|
"total_params": int(total_params),
|
|
"optimizer": optimizer_name,
|
|
"learning_rate": learning_rate,
|
|
"environment": environment,
|
|
},
|
|
)
|
|
for step in range(max_steps):
|
|
window = [samples[(step + offset) % len(samples)] for offset in range(batch_size)]
|
|
queries = [str(sample.inputs[0]["text"]) for sample in window]
|
|
documents = [str(sample.output["positive"]) for sample in window]
|
|
query_batch = encode_retrieval_batch(queries, processor, model.device, max_length)
|
|
document_batch = encode_retrieval_batch(documents, processor, model.device, max_length)
|
|
query_embeddings = pool_embeddings(model(**query_batch), query_batch["attention_mask"])
|
|
document_embeddings = pool_embeddings(model(**document_batch), document_batch["attention_mask"])
|
|
logits = query_embeddings @ document_embeddings.T / 0.05
|
|
labels = torch.arange(logits.shape[0], device=logits.device)
|
|
loss = F.cross_entropy(logits.float(), labels)
|
|
loss.backward()
|
|
grad_l2 = gradient_l2_norm(trainable_tensors)
|
|
before_update = trainable_state_cpu(trainable_tensors)
|
|
optimizer.step()
|
|
update_l2 = trainable_update_l2(before_update, trainable_tensors)
|
|
optimizer_updates.append(
|
|
{
|
|
"step": step + 1,
|
|
"gradient_l2": grad_l2,
|
|
"update_l2": update_l2,
|
|
}
|
|
)
|
|
optimizer.zero_grad(set_to_none=True)
|
|
actual_loss = float(loss.detach().cpu())
|
|
losses.append(actual_loss)
|
|
write_trace(
|
|
trace,
|
|
{
|
|
"event": "step",
|
|
"step": step + 1,
|
|
"sample_ids": [sample.id for sample in window],
|
|
"task": "hme_retrieval_contrastive_pair",
|
|
"loss": actual_loss,
|
|
"optimizer_step": True,
|
|
"gradient_l2": grad_l2,
|
|
"update_l2": update_l2,
|
|
},
|
|
)
|
|
|
|
final_trainable_sha = trainable_state_sha256(trainable_tensors)
|
|
final_trainable_l2 = trainable_l2_norm(trainable_tensors)
|
|
max_update_l2 = max((item["update_l2"] for item in optimizer_updates), default=None)
|
|
loss_delta = losses[0] - losses[-1] if len(losses) >= 2 else None
|
|
model.save_pretrained(output_dir)
|
|
processor.save_pretrained(output_dir)
|
|
write_lora_config(settings, output_dir, TrainingAdapterKind.hme_retrieval, rank, alpha, model_id)
|
|
weight_stats = safetensor_weight_stats(output_dir / "adapter_model.safetensors")
|
|
finished_at = datetime.now(UTC)
|
|
wall_seconds = time.perf_counter() - wall_start
|
|
summary_path = output_dir / "training_summary.json"
|
|
proof_path = output_dir / "training_proof.json"
|
|
training_proof = {
|
|
**contract,
|
|
"adapter_kind": TrainingAdapterKind.hme_retrieval.value,
|
|
"base_model": model_id,
|
|
"dataset_path": str(dataset_path.resolve()),
|
|
"dataset_sha256": sha256_file(dataset_path),
|
|
"sample_count": len(samples),
|
|
"max_steps": max_steps,
|
|
"max_length": max_length,
|
|
"rank": rank,
|
|
"alpha": alpha,
|
|
"target_modules": target_modules,
|
|
"trainable_params": int(trainable_params),
|
|
"total_params": int(total_params),
|
|
"optimizer": optimizer_name,
|
|
"learning_rate": learning_rate,
|
|
"gradient_accumulation_steps": 1,
|
|
"bf16": bool(torch.cuda.is_available()),
|
|
"environment": environment,
|
|
"initial_trainable_sha256": initial_trainable_sha,
|
|
"final_trainable_sha256": final_trainable_sha,
|
|
"initial_trainable_l2": initial_trainable_l2,
|
|
"final_trainable_l2": final_trainable_l2,
|
|
"optimizer_update_count": len(optimizer_updates),
|
|
"max_optimizer_update_l2": max_update_l2,
|
|
"loss_delta": loss_delta,
|
|
"loss_history": losses,
|
|
"optimizer_updates": optimizer_updates,
|
|
"adapter_weight_sha256": weight_stats["sha256"],
|
|
"adapter_weight_stats": weight_stats,
|
|
}
|
|
write_training_summary(proof_path, training_proof)
|
|
write_training_summary(
|
|
summary_path,
|
|
{
|
|
**contract,
|
|
"adapter_kind": TrainingAdapterKind.hme_retrieval.value,
|
|
"base_model": model_id,
|
|
"dataset_path": str(dataset_path.resolve()),
|
|
"dataset_sha256": sha256_file(dataset_path),
|
|
"sample_count": len(samples),
|
|
"max_steps": max_steps,
|
|
"max_length": max_length,
|
|
"rank": rank,
|
|
"alpha": alpha,
|
|
"target_modules": target_modules,
|
|
"trainable_params": int(trainable_params),
|
|
"total_params": int(total_params),
|
|
"optimizer": optimizer_name,
|
|
"learning_rate": learning_rate,
|
|
"gradient_accumulation_steps": 1,
|
|
"bf16": bool(torch.cuda.is_available()),
|
|
"loss_history": losses,
|
|
"final_loss": losses[-1] if losses else None,
|
|
"initial_trainable_sha256": initial_trainable_sha,
|
|
"final_trainable_sha256": final_trainable_sha,
|
|
"initial_trainable_l2": initial_trainable_l2,
|
|
"final_trainable_l2": final_trainable_l2,
|
|
"optimizer_update_count": len(optimizer_updates),
|
|
"max_optimizer_update_l2": max_update_l2,
|
|
"loss_delta": loss_delta,
|
|
"started_at": started_at.isoformat(),
|
|
"finished_at": finished_at.isoformat(),
|
|
"wall_seconds": wall_seconds,
|
|
"environment": environment,
|
|
"peak_cuda_memory_mb": cuda_peak_memory_mb(torch),
|
|
"adapter_weight_sha256": weight_stats["sha256"],
|
|
"adapter_weight_stats": weight_stats,
|
|
},
|
|
)
|
|
with trace_path.open("a", encoding="utf-8", newline="\n") as trace:
|
|
write_trace(
|
|
trace,
|
|
{
|
|
"event": "run_finish",
|
|
**contract,
|
|
"finished_at": finished_at.isoformat(),
|
|
"wall_seconds": wall_seconds,
|
|
"final_loss": losses[-1] if losses else None,
|
|
"initial_trainable_sha256": initial_trainable_sha,
|
|
"final_trainable_sha256": final_trainable_sha,
|
|
"optimizer_update_count": len(optimizer_updates),
|
|
"max_optimizer_update_l2": max_update_l2,
|
|
"loss_delta": loss_delta,
|
|
"peak_cuda_memory_mb": cuda_peak_memory_mb(torch),
|
|
"adapter_weight_sha256": weight_stats["sha256"],
|
|
},
|
|
)
|
|
manifest = AdapterManifest(
|
|
adapter_id=adapter_id(TrainingAdapterKind.hme_retrieval, dataset_path, model_id),
|
|
kind=TrainingAdapterKind.hme_retrieval,
|
|
created_at=datetime.now(UTC),
|
|
proof_version=TRAINING_PROOF_VERSION,
|
|
training_method=TRAINING_METHOD,
|
|
precision=TRAINING_PRECISION,
|
|
framework=TRAINING_FRAMEWORK,
|
|
base_model=model_id,
|
|
dataset_path=str(dataset_path.resolve()),
|
|
dataset_sha256=sha256_file(dataset_path),
|
|
sample_count=len(samples),
|
|
max_steps=max_steps,
|
|
max_length=max_length,
|
|
rank=rank,
|
|
alpha=alpha,
|
|
target_modules=target_modules,
|
|
trainable_params=int(trainable_params),
|
|
total_params=int(total_params),
|
|
optimizer=optimizer_name,
|
|
optimizer_family=optimizer_family(optimizer_name),
|
|
learning_rate=learning_rate,
|
|
gradient_accumulation_steps=1,
|
|
gradient_checkpointing=gradient_checkpointing,
|
|
gradient_checkpointing_mode=GRADIENT_CHECKPOINTING_MODE if gradient_checkpointing else None,
|
|
autocast_dtype=AUTOCAST_DTYPE,
|
|
bf16=torch.cuda.is_available(),
|
|
unsloth_version=package_version("unsloth"),
|
|
triton_version=package_version("triton-windows") or package_version("triton"),
|
|
xformers_version=package_version("xformers"),
|
|
unsloth_model_type=unsloth_model_type,
|
|
final_loss=losses[-1] if losses else None,
|
|
loss_history=losses,
|
|
initial_trainable_sha256=initial_trainable_sha,
|
|
final_trainable_sha256=final_trainable_sha,
|
|
initial_trainable_l2=initial_trainable_l2,
|
|
final_trainable_l2=final_trainable_l2,
|
|
optimizer_update_count=len(optimizer_updates),
|
|
max_optimizer_update_l2=max_update_l2,
|
|
loss_delta=loss_delta,
|
|
training_proof_path=str(proof_path.resolve()),
|
|
training_proof_sha256=sha256_file(proof_path),
|
|
training_trace_path=str(trace_path.resolve()),
|
|
training_trace_sha256=sha256_file(trace_path),
|
|
training_summary_path=str(summary_path.resolve()),
|
|
training_summary_sha256=sha256_file(summary_path),
|
|
training_started_at=started_at,
|
|
training_finished_at=finished_at,
|
|
training_wall_seconds=wall_seconds,
|
|
training_device=str(environment.get("device_name") or environment.get("device") or ""),
|
|
torch_version=str(environment.get("torch_version") or ""),
|
|
cuda_version=(
|
|
str(environment["torch_cuda_version"])
|
|
if environment.get("torch_cuda_version") is not None
|
|
else None
|
|
),
|
|
adapter_weight_sha256=str(weight_stats["sha256"] or ""),
|
|
adapter_tensor_count=int(weight_stats["tensor_count"]),
|
|
adapter_lora_tensor_count=int(weight_stats["lora_tensor_count"]),
|
|
adapter_total_elements=int(weight_stats["total_elements"]),
|
|
adapter_nonzero_elements=int(weight_stats["nonzero_elements"]),
|
|
adapter_max_abs=weight_stats["max_abs"],
|
|
adapter_weights_finite=bool(weight_stats["all_finite"]),
|
|
artifacts=relative_artifacts(output_dir),
|
|
sha256=sha256_tree(output_dir),
|
|
cartridge_id=cartridge_id,
|
|
)
|
|
(output_dir / "adapter_manifest.json").write_text(
|
|
manifest.model_dump_json(indent=2),
|
|
encoding="utf-8",
|
|
)
|
|
return manifest
|