File size: 2,882 Bytes
5cd1287 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 | """Training output, data and optimizer setup shared by the training entrypoint."""
import os
from pathlib import Path
from typing import Tuple
import torch
import torch.distributed as dist
from accelerate.logging import get_logger
from torch.utils.data import DataLoader
from transformers import AutoProcessor, get_scheduler
from starVLA.dataloader import build_dataloader
from starVLA.training.trainer_utils.trainer_tools import build_param_lr_groups
logger = get_logger(__name__)
def load_fast_tokenizer():
return AutoProcessor.from_pretrained("physical-intelligence/fast", trust_remote_code=True)
def setup_directories(cfg) -> Path:
"""Create output directory and checkpoint directory."""
cfg.output_dir = os.path.join(cfg.run_root_dir, cfg.run_id)
output_dir = Path(cfg.output_dir)
if not dist.is_initialized() or dist.get_rank() == 0:
os.makedirs(output_dir, exist_ok=True)
os.makedirs(output_dir / "checkpoints", exist_ok=True)
return output_dir
def prepare_data(cfg, accelerator, output_dir) -> DataLoader:
"""Prepare VLA training data."""
logger.info(f"Creating VLA Dataset with Mixture `{cfg.datasets.vla_data.data_mix}`")
vla_train_dataloader = build_dataloader(cfg=cfg, dataset_py=cfg.datasets.vla_data.dataset_py)
accelerator.dataloader_config.dispatch_batches = False
if dist.is_initialized():
dist.barrier()
return vla_train_dataloader
def setup_optimizer_and_scheduler(model, cfg) -> Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler._LRScheduler]:
"""Set optimizer and scheduler."""
param_groups = build_param_lr_groups(model=model, cfg=cfg)
optimizer_kwargs = {
"lr": cfg.trainer.learning_rate.base,
"betas": tuple(cfg.trainer.optimizer.betas),
"weight_decay": cfg.trainer.optimizer.weight_decay,
"eps": cfg.trainer.optimizer.eps,
}
if str(cfg.trainer.optimizer.name).lower() == "adamw8bit":
import bitsandbytes as bnb
optimizer = bnb.optim.AdamW8bit(param_groups, **optimizer_kwargs)
else:
optimizer = torch.optim.AdamW(
param_groups,
**optimizer_kwargs,
fused=True,
)
if dist.is_initialized() and dist.get_rank() == 0:
for group in optimizer.param_groups:
logger.info(f"LR Group {group['name']}: lr={group['lr']}, num_params={len(group['params'])}")
# Strip keys unknown to transformers' get_scheduler before passing kwargs.
sched_kwargs = {k: v for k, v in cfg.trainer.scheduler_specific_kwargs.items()}
lr_scheduler = get_scheduler(
name=cfg.trainer.lr_scheduler_type,
optimizer=optimizer,
num_warmup_steps=cfg.trainer.num_warmup_steps,
num_training_steps=cfg.trainer.max_train_steps,
scheduler_specific_kwargs=sched_kwargs,
)
return optimizer, lr_scheduler
|