#!/usr/bin/env python3 """Train Qwen2.5-Omni Thinker LoRA + EgoSpeak 10x3 head from AV caches. This is a custom task runner intended to execute inside the ms-swift image. It uses the PEFT library shipped with ms-swift because Swift's normal SFT CLI optimizes generated text tokens, whereas this experiment reads one continuous Thinker hidden state and applies a multi-horizon classification loss. The cached audio and vision towers remain frozen. The reconstructed sequence is forwarded with ``use_cache=False`` so all retained 2-second blocks participate in autograd through the Thinker LoRA adapters. """ from __future__ import annotations import argparse import gc import importlib.metadata import json import math import random import time from pathlib import Path from typing import Any, Sequence import numpy as np import torch from torch import nn from torch.utils.data import DataLoader, Dataset from tqdm.auto import tqdm import streaming_turn_model as core from easycom_head import ( NUM_CLASSES, NUM_HORIZONS, TurnPredictionHead, checkpoint_head_config, egospeak_multihorizon_loss, ) from omni_cached_av import ( SCHEMA_VERSION, CachedBlockStore, collate_cached_sequences, forward_last_state, load_json, pieces_to_cache, reconstruct_open_stream_sequence, save_json, torch_load, ) DEFAULT_LORA_TARGETS = ( "q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj", ) def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--model-id", required=True) parser.add_argument( "--cache-dir", type=Path, action="append", required=True, help="Cache directory from cache_omni_av_pieces.py; repeat for sessions", ) parser.add_argument( "--target", type=Path, help="Override target path for a single --cache-dir", ) parser.add_argument("--output-dir", type=Path, required=True) parser.add_argument("--head-checkpoint", type=Path) parser.add_argument("--max-examples", type=int, default=0) parser.add_argument("--max-steps", type=int, default=200) parser.add_argument("--batch-size", type=int, default=1) parser.add_argument("--gradient-accumulation-steps", type=int, default=1) parser.add_argument("--num-workers", type=int, default=0) parser.add_argument("--block-cache-size", type=int, default=16) parser.add_argument( "--online-current-tail", action="store_true", help=( "Benchmark complete-2s-block caching: read completed history from " "cache but recompute the current 0.2--1.8s AV tail from source media" ), ) parser.add_argument( "--easycom-session-dir", type=Path, help="Original Session_N media; required by --online-current-tail", ) parser.add_argument("--frame-cache-size", type=int, default=64) parser.add_argument("--decoder", choices=("decord", "pynvcodec"), default="decord") parser.add_argument("--decoder-gpu-id", type=int, default=0) parser.add_argument("--decoder-threads", type=int, default=1) parser.add_argument("--max-open-readers", type=int, default=1) parser.add_argument("--learning-rate", type=float, default=2e-5) parser.add_argument("--head-learning-rate", type=float, default=1e-4) parser.add_argument("--weight-decay", type=float, default=0.01) parser.add_argument("--warmup-steps", type=int, default=10) parser.add_argument("--max-grad-norm", type=float, default=1.0) parser.add_argument("--lora-rank", type=int, default=8) parser.add_argument("--lora-alpha", type=int, default=16) parser.add_argument("--lora-dropout", type=float, default=0.05) parser.add_argument( "--lora-target-modules", default=",".join(DEFAULT_LORA_TARGETS), ) parser.add_argument("--head-dropout", type=float, default=0.1) parser.add_argument("--gradient-checkpointing", action="store_true") parser.add_argument("--seed", type=int, default=1337) parser.add_argument("--log-every", type=int, default=5) parser.add_argument("--eval-max-batches", type=int, default=4) parser.add_argument( "--dtype", choices=("bfloat16", "float16"), default="bfloat16", ) parser.add_argument("--device", default="cuda:0") parser.add_argument("--flash-attention-2", action="store_true") parser.add_argument("--overwrite-output", action="store_true") parser.add_argument("--no-progress", action="store_true") return parser.parse_args(argv) def _seed_everything(seed: int) -> None: random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) def _package_version(name: str) -> str | None: try: return importlib.metadata.version(name) except importlib.metadata.PackageNotFoundError: return None def _load_targets(path: Path) -> np.ndarray: targets = np.load(path, mmap_mode="r") if targets.ndim != 2 or targets.shape[1] != NUM_CLASSES: raise ValueError(f"Expected [T,{NUM_CLASSES}] target, got {targets.shape}: {path}") return targets class OmniTurnCacheDataset(Dataset[dict[str, Any]]): """One sample per 200-ms anchor, reconstructed lazily on CPU.""" def __init__( self, cache_dirs: Sequence[Path], *, target_override: Path | None, max_examples: int, seed: int, block_cache_size: int, defer_current_tail: bool = False, ) -> None: if target_override is not None and len(cache_dirs) != 1: raise ValueError("--target can override only one --cache-dir") self.sources: list[dict[str, Any]] = [] self.examples: list[tuple[int, dict[str, Any]]] = [] self.defer_current_tail = bool(defer_current_tail) hidden_size: int | None = None reference_contract: tuple[Any, ...] | None = None for source_index, raw_cache_dir in enumerate(cache_dirs): cache_dir = raw_cache_dir.expanduser().resolve() manifest_path = cache_dir / "manifest.json" manifest = load_json(manifest_path) if int(manifest.get("schema_version", -1)) != SCHEMA_VERSION: raise ValueError( f"Unsupported cache schema in {manifest_path}: " f"{manifest.get('schema_version')}" ) static = torch_load(cache_dir / manifest["static_path"]) target_path = ( target_override.expanduser().resolve() if target_override is not None else Path(manifest["target_path"]) ) targets = _load_targets(target_path) if len(targets) != int(manifest["target_steps"]): raise ValueError( f"Target length changed for {manifest_path}: " f"{len(targets)} vs {manifest['target_steps']}" ) source_hidden = int(manifest["hidden_size"]) hidden_size = source_hidden if hidden_size is None else hidden_size if source_hidden != hidden_size: raise ValueError("Cache directories have different hidden sizes") contract = ( manifest["model_id"], manifest["step_seconds"], manifest["native_block_seconds"], manifest["memory_seconds"], manifest["video_fps"], manifest["positions_per_chunk"], ) reference_contract = contract if reference_contract is None else reference_contract if contract != reference_contract: raise ValueError( "Cache directories do not share the same model/timing contract: " f"{manifest_path}" ) self.sources.append( { "cache_dir": cache_dir, "manifest": manifest, "static": static, "targets": targets, "store": CachedBlockStore(cache_dir, max_blocks=block_cache_size), } ) for anchor in manifest["anchors"]: target_index = int(anchor["target_index"]) if not bool(anchor.get("has_future_target", True)): continue if target_index + NUM_HORIZONS >= len(targets): continue self.examples.append((source_index, anchor)) if not self.examples: raise ValueError("No cache anchors have ten future target steps") if max_examples < 0: raise ValueError("--max-examples must be non-negative") if max_examples and max_examples < len(self.examples): generator = random.Random(seed) selected = generator.sample(range(len(self.examples)), max_examples) self.examples = [self.examples[index] for index in sorted(selected)] self.hidden_size = int(hidden_size or 0) def __len__(self) -> int: return len(self.examples) def __getitem__(self, index: int) -> dict[str, Any]: source_index, anchor = self.examples[index] source = self.sources[source_index] target_index = int(anchor["target_index"]) targets = torch.from_numpy( np.array( source["targets"][ target_index + 1 : target_index + NUM_HORIZONS + 1 ], dtype=np.float32, copy=True, ) ) if targets.shape != (NUM_HORIZONS, NUM_CLASSES): raise RuntimeError( f"Unexpected future target shape {tuple(targets.shape)} at {target_index}" ) item: dict[str, Any] = { "targets": targets, "step_index": target_index, } if self.defer_current_tail: item["source_index"] = source_index item["anchor"] = anchor else: embeddings, positions = reconstruct_open_stream_sequence( store=source["store"], static=source["static"], anchor=anchor, ) item["inputs_embeds"] = embeddings item["position_ids"] = positions return item def _collate_deferred_tail(batch: Sequence[dict[str, Any]]) -> dict[str, Any]: if len(batch) != 1: raise ValueError("--online-current-tail benchmark currently requires batch size 1") return batch[0] class OnlineCurrentTailMaterializer: """Recompute only the current partial block while history stays cached.""" def __init__( self, *, dataset: OmniTurnCacheDataset, model: torch.nn.Module, processor: Any, session_dir: Path, frame_cache_size: int, decoder: str, decoder_gpu_id: int, decoder_threads: int, max_open_readers: int, ) -> None: if len(dataset.sources) != 1: raise ValueError("Online-tail benchmark supports exactly one cache directory") from streaming_turn_model_directly import ( CausalSourceAV, _easycom_directory_specs, ) self.dataset = dataset self.model = model source = dataset.sources[0] manifest = source["manifest"] config = core.StreamConfig( step_seconds=float(manifest["step_seconds"]), native_block_seconds=float(manifest["native_block_seconds"]), memory_seconds=float(manifest["memory_seconds"]), video_fps=float(manifest["video_fps"]), video_min_pixels=int(manifest["video_min_pixels"]), video_max_pixels=int(manifest["video_max_pixels"]), ) media = CausalSourceAV( _easycom_directory_specs(session_dir), output_fps=config.video_fps, required_duration=float(manifest["cached_end_seconds"]), frame_cache_size=frame_cache_size, decoder_backend=decoder, decoder_gpu_id=decoder_gpu_id, decoder_threads=decoder_threads, max_open_readers=max_open_readers, ) self.extractor = core.OpenStreamReplayExtractor( model, processor, config, fixed_grid_av=media ) self.step_seconds = config.step_seconds self.block_seconds = config.native_block_seconds self.steps_per_block = int(source["static"]["steps_per_block"]) self.calls = 0 self.online_av_calls = 0 self.total_seconds = 0.0 self.online_av_seconds = 0.0 def reset_timing(self) -> None: self.calls = 0 self.online_av_calls = 0 self.total_seconds = 0.0 self.online_av_seconds = 0.0 @torch.no_grad() def __call__(self, item: dict[str, Any]) -> dict[str, Any]: started = time.perf_counter() source_index = int(item["source_index"]) source = self.dataset.sources[source_index] anchor = item["anchor"] tail_steps = int(anchor["tail_steps"]) current_piece = None if tail_steps < self.steps_per_block: av_started = time.perf_counter() block_start = int(anchor["block_index"]) * self.block_seconds interval_end = block_start + tail_steps * self.step_seconds pieces = self.extractor._prepare_pieces( video_path=None, audio_path=None, interval=(block_start, interval_end), use_audio_in_video=True, ) closure_embeddings = self.model.get_input_embeddings()( pieces.closure_ids ) current_piece = pieces_to_cache( pieces, closure_embeddings=closure_embeddings, floating_dtype=torch.bfloat16, ) self.online_av_seconds += time.perf_counter() - av_started self.online_av_calls += 1 embeddings, positions = reconstruct_open_stream_sequence( store=source["store"], static=source["static"], anchor=anchor, current_piece=current_piece, ) batch = collate_cached_sequences( [ { "inputs_embeds": embeddings, "position_ids": positions, "targets": item["targets"], "step_index": item["step_index"], } ] ) self.total_seconds += time.perf_counter() - started self.calls += 1 return batch def _load_initial_head( checkpoint: Path | None, *, hidden_size: int, dropout: float, ) -> TurnPredictionHead: if checkpoint is None: return TurnPredictionHead(hidden_size=hidden_size, dropout=dropout) try: payload = torch.load(checkpoint, map_location="cpu", weights_only=True) except (TypeError, RuntimeError): payload = torch.load(checkpoint, map_location="cpu") if not isinstance(payload, dict): raise ValueError(f"Expected dict head checkpoint: {checkpoint}") config = checkpoint_head_config(payload) if int(config["hidden_size"]) != hidden_size: raise ValueError( f"Head hidden size {config['hidden_size']} != cache hidden size {hidden_size}" ) head = TurnPredictionHead( hidden_size=hidden_size, num_horizons=int(config["num_horizons"]), num_classes=int(config["num_classes"]), dropout=float(config["dropout"]), ) head.load_state_dict(payload.get("head", payload)) print(f"Initialized classification head from {checkpoint}") return head def _prepare_language_lora( model: torch.nn.Module, *, rank: int, alpha: int, dropout: float, target_modules: Sequence[str], gradient_checkpointing: bool, ) -> torch.nn.Module: try: from peft import LoraConfig, get_peft_model except ImportError as exc: raise RuntimeError( "PEFT is required. The ms-swift 3.12.5 image should include it." ) from exc model.requires_grad_(False) language_model = model.model matched = [ name for name, module in language_model.named_modules() if isinstance(module, nn.Linear) and name.rsplit(".", 1)[-1] in target_modules ] if not matched: raise RuntimeError( "No language-model linear layers matched LoRA targets " f"{list(target_modules)}" ) config = LoraConfig( r=rank, lora_alpha=alpha, lora_dropout=dropout, target_modules=list(target_modules), bias="none", ) language_model = get_peft_model(language_model, config) language_model.config.use_cache = False if gradient_checkpointing: try: language_model.gradient_checkpointing_enable( gradient_checkpointing_kwargs={"use_reentrant": False} ) except TypeError: language_model.gradient_checkpointing_enable() language_model.print_trainable_parameters() print( f"LoRA matched {len(matched)} language layers; first={matched[0]}, " f"last={matched[-1]}" ) return language_model def _lr_lambda(step: int, *, total_steps: int, warmup_steps: int) -> float: if warmup_steps > 0 and step < warmup_steps: return float(step + 1) / float(warmup_steps) remaining = max(1, total_steps - warmup_steps) progress = min(1.0, max(0.0, (step - warmup_steps) / remaining)) return 0.5 * (1.0 + math.cos(math.pi * progress)) @torch.inference_mode() def _evaluate_loss( language_model: torch.nn.Module, head: TurnPredictionHead, loader: DataLoader, *, device: torch.device, model_dtype: torch.dtype, max_batches: int, materialize_batch: Any | None = None, ) -> float: language_was_training = language_model.training head_was_training = head.training language_model.eval() head.eval() losses: list[float] = [] for batch_index, batch in enumerate(loader): if max_batches > 0 and batch_index >= max_batches: break if materialize_batch is not None: batch = materialize_batch(batch) inputs_embeds = batch["inputs_embeds"].to( device=device, dtype=model_dtype, non_blocking=True ) position_ids = batch["position_ids"].to(device=device, non_blocking=True) attention_mask = batch["attention_mask"].to(device=device, non_blocking=True) readout_indices = batch["readout_indices"].to(device=device, non_blocking=True) targets = batch["targets"].to(device=device, dtype=torch.float32) states = forward_last_state( language_model, inputs_embeds=inputs_embeds, position_ids=position_ids, attention_mask=attention_mask, readout_indices=readout_indices, ) logits = head(states.float()) losses.append(float(egospeak_multihorizon_loss(logits, targets).item())) language_model.train(language_was_training) head.train(head_was_training) return float(np.mean(losses)) def _jsonable_args(args: argparse.Namespace) -> dict[str, Any]: output: dict[str, Any] = {} for key, value in vars(args).items(): if isinstance(value, Path): output[key] = str(value) elif isinstance(value, list) and value and isinstance(value[0], Path): output[key] = [str(item) for item in value] else: output[key] = value return output def main(argv: Sequence[str] | None = None) -> None: args = parse_args(argv) if args.max_steps <= 0: raise ValueError("--max-steps must be positive") if args.batch_size <= 0 or args.gradient_accumulation_steps <= 0: raise ValueError("Batch size and gradient accumulation must be positive") if args.online_current_tail: if args.easycom_session_dir is None: raise ValueError( "--easycom-session-dir is required by --online-current-tail" ) if len(args.cache_dir) != 1: raise ValueError("Online-tail benchmark requires exactly one --cache-dir") if args.batch_size != 1: raise ValueError("Online-tail benchmark currently requires --batch-size 1") if args.num_workers != 0: raise ValueError("Online-tail benchmark currently requires --num-workers 0") if args.num_workers != 0: print( "Warning: each worker owns an independent block-file LRU; use 0 for " "the 16-example pilot and increase only after profiling." ) output_dir = args.output_dir.expanduser().resolve() known_outputs = ( output_dir / "adapter_config.json", output_dir / "adapter", output_dir / "turn_head.pt", output_dir / "train_summary.json", ) if not args.overwrite_output and any(path.exists() for path in known_outputs): raise FileExistsError( f"Training output already exists in {output_dir}; pass --overwrite-output" ) output_dir.mkdir(parents=True, exist_ok=True) _seed_everything(args.seed) core._configure_qwen_logging() device = torch.device(args.device) if device.type != "cuda": raise ValueError("This LoRA pilot is intended for one CUDA GPU") dataset = OmniTurnCacheDataset( args.cache_dir, target_override=args.target, max_examples=args.max_examples, seed=args.seed, block_cache_size=args.block_cache_size, defer_current_tail=args.online_current_tail, ) collate_fn = ( _collate_deferred_tail if args.online_current_tail else collate_cached_sequences ) generator = torch.Generator().manual_seed(args.seed) train_loader = DataLoader( dataset, batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers, pin_memory=True, persistent_workers=args.num_workers > 0, collate_fn=collate_fn, generator=generator, ) eval_loader = DataLoader( dataset, batch_size=args.batch_size, shuffle=False, num_workers=0, pin_memory=True, collate_fn=collate_fn, ) print( f"Cached training examples: {len(dataset)}, hidden_size={dataset.hidden_size}, " f"batch={args.batch_size} x accum={args.gradient_accumulation_steps}" ) base_model, processor = core.load_qwen_omni( args.model_id, dtype=args.dtype, device_map=args.device, use_flash_attention_2=args.flash_attention_2, ) target_modules = tuple( value.strip() for value in args.lora_target_modules.split(",") if value.strip() ) if not target_modules: raise ValueError("--lora-target-modules is empty") language_model = _prepare_language_lora( base_model, rank=args.lora_rank, alpha=args.lora_alpha, dropout=args.lora_dropout, target_modules=target_modules, gradient_checkpointing=args.gradient_checkpointing, ) materialize_batch: OnlineCurrentTailMaterializer | None = None if args.online_current_tail: assert args.easycom_session_dir is not None materialize_batch = OnlineCurrentTailMaterializer( dataset=dataset, model=base_model, processor=processor, session_dir=args.easycom_session_dir, frame_cache_size=args.frame_cache_size, decoder=args.decoder, decoder_gpu_id=args.decoder_gpu_id, decoder_threads=args.decoder_threads, max_open_readers=args.max_open_readers, ) print( "Cache policy benchmark: complete 2s history cached; " "current 0.2--1.8s tail computed online" ) else: # All AV phases are cached, so the modality towers can be discarded. del base_model del processor gc.collect() torch.cuda.empty_cache() model_dtype = next(language_model.parameters()).dtype head = _load_initial_head( args.head_checkpoint, hidden_size=dataset.hidden_size, dropout=args.head_dropout, ).to(device=device, dtype=torch.float32) lora_parameters = [parameter for parameter in language_model.parameters() if parameter.requires_grad] if not lora_parameters: raise RuntimeError("LoRA preparation produced no trainable parameters") optimizer = torch.optim.AdamW( [ {"params": lora_parameters, "lr": args.learning_rate}, {"params": head.parameters(), "lr": args.head_learning_rate}, ], weight_decay=args.weight_decay, ) scheduler = torch.optim.lr_scheduler.LambdaLR( optimizer, lr_lambda=lambda step: _lr_lambda( step, total_steps=args.max_steps, warmup_steps=args.warmup_steps ), ) initial_loss = _evaluate_loss( language_model, head, eval_loader, device=device, model_dtype=model_dtype, max_batches=args.eval_max_batches, materialize_batch=materialize_batch, ) print(f"Initial cached loss: {initial_loss:.6f}") if materialize_batch is not None: materialize_batch.reset_timing() language_model.train() head.train() optimizer.zero_grad(set_to_none=True) iterator = iter(train_loader) history: list[dict[str, float | int]] = [] progress = tqdm( total=args.max_steps, desc="Thinker LoRA", unit="update", disable=args.no_progress, ) training_started = time.perf_counter() for update_step in range(1, args.max_steps + 1): accumulated_loss = 0.0 for _ in range(args.gradient_accumulation_steps): try: batch = next(iterator) except StopIteration: iterator = iter(train_loader) batch = next(iterator) if materialize_batch is not None: batch = materialize_batch(batch) inputs_embeds = batch["inputs_embeds"].to( device=device, dtype=model_dtype, non_blocking=True ) position_ids = batch["position_ids"].to(device=device, non_blocking=True) attention_mask = batch["attention_mask"].to(device=device, non_blocking=True) readout_indices = batch["readout_indices"].to(device=device, non_blocking=True) targets = batch["targets"].to(device=device, dtype=torch.float32) states = forward_last_state( language_model, inputs_embeds=inputs_embeds, position_ids=position_ids, attention_mask=attention_mask, readout_indices=readout_indices, ) logits = head(states.float()) loss = egospeak_multihorizon_loss(logits, targets) (loss / args.gradient_accumulation_steps).backward() accumulated_loss += float(loss.detach()) / args.gradient_accumulation_steps trainable_parameters = [*lora_parameters, *head.parameters()] grad_norm = torch.nn.utils.clip_grad_norm_( trainable_parameters, args.max_grad_norm ) optimizer.step() scheduler.step() optimizer.zero_grad(set_to_none=True) record = { "step": update_step, "loss": accumulated_loss, "grad_norm": float(grad_norm), "lora_lr": float(optimizer.param_groups[0]["lr"]), "head_lr": float(optimizer.param_groups[1]["lr"]), } history.append(record) progress.update(1) progress.set_postfix( loss=f"{accumulated_loss:.4f}", lr=f"{optimizer.param_groups[0]['lr']:.2g}", ) if update_step == 1 or update_step % args.log_every == 0: print( f"step={update_step} loss={accumulated_loss:.6f} " f"grad_norm={float(grad_norm):.4f} " f"lora_lr={optimizer.param_groups[0]['lr']:.3g}", flush=True, ) progress.close() training_seconds = time.perf_counter() - training_started online_timing = None if materialize_batch is not None: online_timing = { "materialize_calls": materialize_batch.calls, "materialize_seconds": materialize_batch.total_seconds, "online_av_calls": materialize_batch.online_av_calls, "online_av_seconds": materialize_batch.online_av_seconds, } final_loss = _evaluate_loss( language_model, head, eval_loader, device=device, model_dtype=model_dtype, max_batches=args.eval_max_batches, materialize_batch=materialize_batch, ) adapter_dir = output_dir / "adapter" language_model.save_pretrained(adapter_dir) torch.save( { "head": head.state_dict(), "head_config": { "hidden_size": dataset.hidden_size, "num_horizons": NUM_HORIZONS, "num_classes": NUM_CLASSES, "dropout": head.dropout, }, "initial_head_checkpoint": ( str(args.head_checkpoint.resolve()) if args.head_checkpoint else None ), "training_step": args.max_steps, }, output_dir / "turn_head.pt", ) summary = { "schema_version": 1, "status": "pilot" if args.max_examples else "training", "model_id": args.model_id, "cache_dirs": [str(path.expanduser().resolve()) for path in args.cache_dir], "num_examples": len(dataset), "initial_loss": initial_loss, "final_loss": final_loss, "training_seconds": training_seconds, "updates_per_second": args.max_steps / training_seconds, "online_tail_timing": online_timing, "loss_reduction_ratio": ( (initial_loss - final_loss) / initial_loss if initial_loss else None ), "peak_cuda_allocated_gib": torch.cuda.max_memory_allocated(device) / (1024 ** 3), "packages": { "torch": torch.__version__, "transformers": _package_version("transformers"), "peft": _package_version("peft"), "ms-swift": _package_version("ms-swift"), }, "args": _jsonable_args(args), "history": history, } save_json(output_dir / "train_summary.json", summary) print(f"Final cached loss: {final_loss:.6f} (initial {initial_loss:.6f})") print( f"Training speed: {training_seconds:.3f}s / {args.max_steps} updates = " f"{args.max_steps / training_seconds:.3f} update/s" ) if online_timing is not None: print( "Online current-tail AV: " f"{online_timing['online_av_seconds']:.3f}s / " f"{online_timing['online_av_calls']} calls" ) print(f"Saved LoRA adapter to {adapter_dir}") print(f"Saved turn head to {output_dir / 'turn_head.pt'}") print(f"Peak CUDA allocated: {summary['peak_cuda_allocated_gib']:.2f} GiB") if __name__ == "__main__": main()