| """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'])}") |
|
|
| |
| 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 |
|
|
