Buckets:
| """UniAR processor, LoRA model, and optimizer construction.""" | |
| from __future__ import annotations | |
| from pathlib import Path | |
| import torch | |
| from transformers import AutoProcessor | |
| from uniar import UniARForConditionalGeneration, configure_processor | |
| def unwrap_ar_model(model): | |
| return model.get_base_model() if hasattr(model, "get_base_model") else model | |
| def model_autocast(model): | |
| """Use the BF16 CUDA execution contract while retaining FP32 master weights.""" | |
| device = next(unwrap_ar_model(model).parameters()).device | |
| return torch.autocast( | |
| device_type=device.type, | |
| dtype=torch.bfloat16, | |
| enabled=device.type == "cuda", | |
| ) | |
| def load_processor(model_path: str | Path): | |
| return configure_processor(AutoProcessor.from_pretrained(model_path, use_fast=True)) | |
| def validate_lora_adapter_config(actual, expected) -> None: | |
| """Fail when a loaded PEFT adapter contradicts the recorded run config.""" | |
| def enum_value(value): | |
| return getattr(value, "value", value) | |
| comparisons = { | |
| "rank": (getattr(actual, "r", None), expected.rank), | |
| "alpha": (getattr(actual, "lora_alpha", None), expected.alpha), | |
| "dropout": (getattr(actual, "lora_dropout", None), expected.dropout), | |
| "use_rslora": (getattr(actual, "use_rslora", None), expected.use_rslora), | |
| "targets": ( | |
| set(getattr(actual, "target_modules", ()) or ()), | |
| set(expected.targets), | |
| ), | |
| "modules_to_save": ( | |
| set(getattr(actual, "modules_to_save", ()) or ()), | |
| set(expected.modules_to_save), | |
| ), | |
| "rank_pattern": (getattr(actual, "rank_pattern", {}) or {}, {}), | |
| "alpha_pattern": (getattr(actual, "alpha_pattern", {}) or {}, {}), | |
| "peft_type": (enum_value(getattr(actual, "peft_type", None)), "LORA"), | |
| "task_type": ( | |
| enum_value(getattr(actual, "task_type", None)), | |
| "CAUSAL_LM", | |
| ), | |
| "bias": (getattr(actual, "bias", None), "none"), | |
| "use_dora": (getattr(actual, "use_dora", None), False), | |
| "lora_bias": (getattr(actual, "lora_bias", None), False), | |
| } | |
| mismatches = [ | |
| f"{name}: loaded={loaded!r}, configured={configured!r}" | |
| for name, (loaded, configured) in comparisons.items() | |
| if loaded != configured | |
| ] | |
| if mismatches: | |
| raise ValueError("loaded LoRA adapter config mismatch: " + "; ".join(mismatches)) | |
| def _validate_model_lora_adapter(model, expected) -> None: | |
| adapters = getattr(model, "peft_config", None) | |
| if not isinstance(adapters, dict) or len(adapters) != 1: | |
| raise ValueError("training requires exactly one loaded LoRA adapter config") | |
| validate_lora_adapter_config(next(iter(adapters.values())), expected) | |
| def load_model(config, device: torch.device, *, for_training: bool): | |
| """Load the owned checkpoint and attach the configured LoRA adapter.""" | |
| base = UniARForConditionalGeneration.from_pretrained( | |
| config.model.path, | |
| dtype={"bfloat16": torch.bfloat16}[config.model.dtype], | |
| attn_implementation=config.model.attention, | |
| low_cpu_mem_usage=True, | |
| ).to(device) | |
| base.config.use_cache = False | |
| if config.model.freeze_vision_encoder: | |
| base.model.visual.requires_grad_(False) | |
| base.model.visual.eval() | |
| resume = Path(config.checkpoint.resume_from) if config.checkpoint.resume_from else None | |
| initial_adapter = getattr(config.model, "initial_adapter", None) | |
| adapter_path = ( | |
| resume / "adapter" | |
| if resume | |
| else Path(initial_adapter) | |
| if initial_adapter | |
| else None | |
| ) | |
| if adapter_path is not None and not adapter_path.exists(): | |
| raise FileNotFoundError(f"model adapter is missing: {adapter_path}") | |
| model = base | |
| if for_training: | |
| from peft import LoraConfig, PeftModel, TaskType, get_peft_model | |
| if adapter_path is not None: | |
| model = PeftModel.from_pretrained(base, adapter_path, is_trainable=True) | |
| else: | |
| model = get_peft_model( | |
| base, | |
| LoraConfig( | |
| task_type=TaskType.CAUSAL_LM, | |
| r=config.lora.rank, | |
| lora_alpha=config.lora.alpha, | |
| lora_dropout=config.lora.dropout, | |
| bias="none", | |
| use_rslora=config.lora.use_rslora, | |
| target_modules=config.lora.targets, | |
| modules_to_save=config.lora.modules_to_save, | |
| ), | |
| ) | |
| elif adapter_path is not None: | |
| from peft import PeftModel | |
| model = PeftModel.from_pretrained(base, adapter_path, is_trainable=False) | |
| if adapter_path is not None or for_training: | |
| _validate_model_lora_adapter(model, config.lora) | |
| ar_model = unwrap_ar_model(model) | |
| ar_model.config.gradient_checkpointing_use_reentrant = ( | |
| config.model.gradient_checkpointing_use_reentrant | |
| ) | |
| if config.model.freeze_vision_encoder: | |
| ar_model.model.visual.requires_grad_(False) | |
| ar_model.model.visual.eval() | |
| model.config.use_cache = False | |
| if for_training and config.model.gradient_checkpointing: | |
| model.gradient_checkpointing_enable( | |
| gradient_checkpointing_kwargs={ | |
| "use_reentrant": config.model.gradient_checkpointing_use_reentrant | |
| } | |
| ) | |
| model.enable_input_require_grads() | |
| if for_training and config.optimizer.visual_head_fp32: | |
| for parameter in ar_model.output_layer_vistok.parameters(): | |
| if parameter.requires_grad: | |
| parameter.data = parameter.data.float() | |
| return model | |
| def trainable_parameter_groups( | |
| model, | |
| ) -> tuple[list[torch.nn.Parameter], list[torch.nn.Parameter]]: | |
| ar_model = unwrap_ar_model(model) | |
| head_ids = { | |
| id(parameter) | |
| for parameter in ar_model.output_layer_vistok.parameters() | |
| if parameter.requires_grad | |
| } | |
| visual_head = [] | |
| backbone = [] | |
| for parameter in model.parameters(): | |
| if parameter.requires_grad: | |
| (visual_head if id(parameter) in head_ids else backbone).append(parameter) | |
| if not backbone or not visual_head: | |
| raise RuntimeError( | |
| "expected non-empty backbone and visual-head trainable parameter groups: " | |
| f"backbone={len(backbone)}, visual_head={len(visual_head)}" | |
| ) | |
| return backbone, visual_head | |
| def build_optimizer(model, config): | |
| backbone, visual_head = trainable_parameter_groups(model) | |
| optimizer = torch.optim.AdamW( | |
| [ | |
| {"params": backbone, "lr": config.optimizer.learning_rate, "name": "backbone"}, | |
| { | |
| "params": visual_head, | |
| "lr": config.optimizer.visual_head_learning_rate, | |
| "name": "visual_head", | |
| }, | |
| ], | |
| betas=(config.optimizer.beta1, config.optimizer.beta2), | |
| eps=config.optimizer.epsilon, | |
| weight_decay=config.optimizer.weight_decay, | |
| fused=True, | |
| ) | |
| return optimizer, backbone, visual_head | |
| __all__ = [ | |
| "build_optimizer", | |
| "load_model", | |
| "load_processor", | |
| "model_autocast", | |
| "trainable_parameter_groups", | |
| "unwrap_ar_model", | |
| "validate_lora_adapter_config", | |
| ] | |
Xet Storage Details
- Size:
- 7.29 kB
- Xet hash:
- 4ee074934f11f0c9a772716c05f1426b0ef25ba3dd73211e000fdc1daaed606c
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.