학습 가이드¶

LORA_EXCLUDED_PATHS = ("switch_mlp", "lm_head")
LORA_EXCLUDED_NAMES = ("x_proj", "dt_proj", "gate", "shared_expert_gate")

from bit_axon.training.lora import apply_lora_to_model

# DoRA 적용 (기본값): 래핑된 레이어 경로 목록 반환
wrapped = apply_lora_to_model(model, rank=8, dropout=0.0, scale=20.0, use_dora=True)

# 일반 LoRA 적용
wrapped = apply_lora_to_model(model, rank=8, dropout=0.0, scale=20.0, use_dora=False)

from bit_axon.training.merging import merge_adapters

# LoRA 퓨즈: W_fused = W + scale · B^T · A^T
# DoRA 퓨즈: W_fused = (m / ||W + delta||) · (W + delta)
merge_adapters(model)

from bit_axon.training.packing import SequencePacker

packer = SequencePacker(max_seq_len=2048, eos_token_id=151645)

for token_ids, loss_mask in dataset:
    packed_batches = packer.add_example(token_ids, loss_mask)
    for batch in packed_batches:
        process(batch)  # PackedBatch(token_ids, loss_mask)

final = packer.flush()  # 남은 토큰을 max_seq_len으로 패딩

from bit_axon.training.cooling import CoolingScheduler, ThermalPolicy, ThermalShutdownError

policy = ThermalPolicy(
    max_speed_temp=75.0,   # 배치 축소 구간 시작
    pause_temp=85.0,       # 이 온도 이상에서 학습 일시 정지
    stop_temp=95.0,        # 이 온도 이상에서 학습 중지
    pause_duration=0.5,    # 일시 정지 중 대기 간격 (초)
)

cooling = CoolingScheduler(monitor, policy)

# 각 학습 스텝 전에 호출:
cooling.check_before_step(step)  # 일시 정지 또는 ThermalShutdownError 발생

# 온도 대기로 소요된 총 시간 확인:
print(f"Paused {cooling.total_pause_time:.1f}s for cooling")

from bit_axon.training.config import TrainingConfig
from bit_axon.training.data import ORPODataset
from bit_axon.training.orpo_trainer import ORPOTrainer

config = TrainingConfig(
    training_mode="orpo",
    beta=0.1,
    max_steps=2000,
)

dataset = ORPODataset("prefs.jsonl", tokenizer, max_seq_len=2048)
trainer = ORPOTrainer(model, config, dataset, cooling_scheduler=cooling)
result = trainer.train()

# Result keys:
# step, loss, grad_norm, chosen_reward, rejected_reward, reward_margin, reward_accuracy

from bit_axon.training.config import TrainingConfig

config = TrainingConfig(
    # 옵티마이저
    learning_rate=1e-4,       # 웜업 후 최대 LR
    weight_decay=0.01,        # AdamW 가중치 감쇠
    warmup_steps=100,         # 선형 웜업 스텝
    max_steps=10_000,         # 전체 학습 스텝
    max_grad_norm=1.0,        # 기울기 클리핑 임계값
    grad_accum_steps=4,       # 기울기 누적 스텝

    # LoRA / DoRA
    lora_rank=8,              # 저랭크 분해 랭크
    lora_dropout=0.0,         # 어댑터 경로의 dropout
    lora_scale=20.0,          # 어댑터 출력 스케일링
    use_dora=True,            # DoRA 사용 (가중치 분해 LoRA)

    # ORPO 정렬
    beta=0.1,                 # ORPO 선호도 강도
    training_mode="sft",      # "sft" 또는 "orpo"

    # 양자화
    quantize_bits=4,          # 베이스 가중치 비트 폭 (NF4)
    quantize_group_size=64,   # 양자화 그룹 크기

    # 데이터
    batch_size=1,             # 배치당 시퀀스 수
    max_seq_len=2048,         # 패킹 대상 길이

    # 체크포인팅
    save_every=500,           # N 스텝마다 체크포인트 저장
    eval_every=500,           # N 스텝마다 평가
    output_dir="checkpoints", # 체크포인트 디렉토리

    # 온도 임계값 (°C)
    temp_max_speed=75.0,      # 배치 축소 구간
    temp_pause=85.0,          # 학습 일시 정지
    temp_stop=95.0,           # 학습 중지
    temp_poll_interval=1.0,   # 온도 폴링 간격 (초)

    # 기타
    seed=42,                  # 난수 시드
)

import mlx.core as mx
from bit_axon import BitAxonConfig, BitAxonModel
from bit_axon.tokenizer import QwenTokenizerWrapper
from bit_axon.training import TrainingConfig, Trainer, apply_lora_to_model
from bit_axon.training.data import SFTDataset, CacheDataset
from bit_axon.training.cooling import CoolingScheduler, ThermalPolicy

# 1. 모델 로드
model_config = BitAxonConfig()
model = BitAxonModel(model_config)

# 2. DoRA 어댑터 적용
wrapped_layers = apply_lora_to_model(
    model,
    rank=8,
    dropout=0.0,
    scale=20.0,
    use_dora=True,
)
mx.eval(model.parameters())

# 3. 데이터 준비
tokenizer = QwenTokenizerWrapper("Qwen/Qwen2.5-3B")
dataset = CacheDataset(SFTDataset("data.json", tokenizer, max_seq_len=2048))
val_dataset = SFTDataset("val.json", tokenizer, max_seq_len=2048)

# 4. 학습 설정
config = TrainingConfig(
    learning_rate=1e-4,
    max_steps=5000,
    grad_accum_steps=4,
    save_every=500,
    eval_every=500,
    output_dir="checkpoints/my-run",
)

# 5. 온도 모니터링 설정
policy = ThermalPolicy(pause_temp=85.0, stop_temp=95.0)
cooling = CoolingScheduler(thermal_monitor, policy)

# 6. 학습
trainer = Trainer(model, config, dataset, val_dataset, cooling_scheduler=cooling)
result = trainer.train()

print(f"Step {result['step']}: loss={result['loss']:.4f}, grad_norm={result['grad_norm']:.4f}")

from bit_axon.training.config import TrainingConfig
from bit_axon.training.data import ORPODataset
from bit_axon.training.orpo_trainer import ORPOTrainer

config = TrainingConfig(
    training_mode="orpo",
    beta=0.1,
    max_steps=2000,
    save_every=500,
)

dataset = ORPODataset("prefs.jsonl", tokenizer, max_seq_len=2048)
trainer = ORPOTrainer(model, config, dataset, cooling_scheduler=cooling)
result = trainer.train()

print(f"Reward margin: {result['reward_margin']:.4f}")
print(f"Reward accuracy: {result['reward_accuracy']:.2f}")

from bit_axon.training.lora import apply_lora_to_model, LoRALinear
from bit_axon.training.dora import DoRALinear

# 모든 대상 레이어에 적용 (래핑된 경로 목록 반환)
wrapped = apply_lora_to_model(model, rank=8, scale=20.0, use_dora=True)
print(f"Wrapped {len(wrapped)} layers: {wrapped[:3]}...")

# 개별 레이어 제어
dora_layer = DoRALinear.from_base(existing_linear, r=8, scale=20.0)
lora_layer = LoRALinear.from_base(existing_linear, r=8, scale=20.0)

# 어댑터를 베이스 가중치에 퓨즈
fused_linear = dora_layer.fuse()  # 크기 벡터로 재정규화

from bit_axon.training.checkpoint import (
    save_checkpoint,
    load_checkpoint,
    get_latest_checkpoint,
    save_adapter_only,
)

# 수동으로 체크포인트 저장
ckpt_path = save_checkpoint(model, optimizer, step=1500, loss=1.23, output_dir="checkpoints")

# 최신 체크포인트 찾기
latest = get_latest_checkpoint("checkpoints")  # Path 또는 None 반환

# 체크포인트 로드 (모델 가중치 + 옵티마이저 상태 복원)
step, loss = load_checkpoint(model, optimizer, latest)

# 어댑터 가중치만 내보내기 (공유 또는 배포용)
save_adapter_only(model, "my_adapter.safetensors")

from bit_axon.training.merging import merge_adapters, save_merged_model, load_and_merge

# 원스텝: 베이스 + 어댑터 로드, 병합, 재양자화, 저장
load_and_merge(
    base_model_path="./model",
    adapter_path="./checkpoints/final_adapter.safetensors",
    output_dir="./merged-model",
    quantize_after_merge=True,
    bits=4,
    group_size=64,
    lora_rank=8,
)

# 수동 단계별 실행
merge_adapters(model)
save_merged_model(model, "./merged-model", config=model_config)

from bit_axon.training.data import SFTDataset, CacheDataset

raw = SFTDataset("train.jsonl", tokenizer, max_seq_len=2048)
cached = CacheDataset(raw)

# 첫 접근 시 토큰화 및 캐싱; 이후 접근은 캐시에서 로드
for token_ids, loss_mask in cached:
    train_step(token_ids, loss_mask)