Spaces:
Running on L40S
Running on L40S
File size: 2,690 Bytes
9f818c5 | 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 | # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: OpenMDW-1.1
import torch
import wandb
import cosmos_framework.data.vfm.sequence_packing as sequence_packing
from cosmos_framework.callbacks.every_n import EveryN
from cosmos_framework.model._base import ImaginaireModel
from cosmos_framework.trainer import ImaginaireTrainer
class SequencePackingPadding(EveryN):
"""
Callback that saves lengths to which und and gen sequences are padded. This information will be used
to compute FLOPs done during training.
Args:
every_n (int): Frequency with which callback is run during training.
"""
def __init__(self, every_n: int = 500):
super().__init__(every_n=every_n, step_size=1, barrier_after_run=False, run_at_start=True)
def every_n_impl(
self,
trainer: ImaginaireTrainer,
model: ImaginaireModel,
data_batch: dict[str, torch.Tensor],
output_batch: dict[str, torch.Tensor],
loss: torch.Tensor,
iteration: int,
) -> None:
if wandb.run:
log_dict = {
"SequencePackingPadding/max_causal_len_image_batch": sequence_packing.MAX_CAUSAL_LEN_IMAGE_BATCH,
"SequencePackingPadding/max_full_len_image_batch": sequence_packing.MAX_FULL_LEN_IMAGE_BATCH,
"SequencePackingPadding/max_causal_len_video_batch": sequence_packing.MAX_CAUSAL_LEN_VIDEO_BATCH,
"SequencePackingPadding/max_full_len_video_batch": sequence_packing.MAX_FULL_LEN_VIDEO_BATCH,
}
modality = "video"
if "is_image_batch" in output_batch:
modality = "image" if output_batch["is_image_batch"] else "video"
if "und_token_length" in output_batch:
log_dict[f"SequencePackingPadding/und_token_length_{modality}"] = output_batch["und_token_length"]
if "gen_token_length" in output_batch:
log_dict[f"SequencePackingPadding/gen_token_length_{modality}"] = output_batch["gen_token_length"]
if "action_token_length" in output_batch:
log_dict[f"SequencePackingPadding/action_token_length"] = output_batch["action_token_length"]
if "sound_token_length" in output_batch:
log_dict[f"SequencePackingPadding/sound_token_length"] = output_batch["sound_token_length"]
if "vision_token_length" in output_batch:
log_dict[f"SequencePackingPadding/vision_token_length"] = output_batch["vision_token_length"]
wandb.log(
log_dict,
step=iteration,
)
|