Add some logging.

This commit is contained in:
Davids048
2026-08-21 08:33:09 +00:00
parent 0980ca563f
commit 628591b620
3 changed files with 23 additions and 2 deletions
+6 -1
View File
@@ -19,7 +19,12 @@ from fastvideo.attention.backends.abstract import (
from fastvideo.logger import init_logger
logger = init_logger(__name__)
logger.info("Using FlashAttention-%s backend", fa_version)
# Every worker records the loaded FlashAttention implementation so a
# distributed profiling log contains one backend receipt per rank.
logger.info("Worker %s Using FlashAttention-%s backend",
os.environ.get("RANK", "0"),
fa_version,
local_main_process_only=False)
# FP4 FA4 support: quantize Q/K to NVFP4 E2M1 for block-scaled MMA on Blackwell.
# Requires: flash-attention-fp4, flashinfer, cutlass-dsl. Enable via nvfp4_fa4=True kwarg.
+6 -1
View File
@@ -1057,7 +1057,12 @@ class TransformerLoader(ComponentLoader):
# so recording here makes the decision readable from the loaded
# transformer — and records the narrowed one for teacher/critic.
resolved = record_resolved_attention_backend(dit_config)
logger.info("transformer attention backend: %s", resolved.name if resolved else "automatic selection")
# Every worker records its resolved backend so distributed profile
# snapshots can prove that all ranks use the requested kernels.
logger.info("Worker %s transformer attention backend: %s",
os.environ.get("RANK", "0"),
resolved.name if resolved else "automatic selection",
local_main_process_only=False)
model = maybe_load_fsdp_model(
model_cls=model_cls,
init_params={
+11
View File
@@ -4,6 +4,7 @@ from typing import Any, cast
import torch
import fastvideo.envs as envs
from fastvideo.distributed import (cleanup_dist_env_and_memory, maybe_init_distributed_environment_and_model_parallel)
from fastvideo.distributed.parallel_state import get_local_torch_device
from fastvideo.fastvideo_args import FastVideoArgs
@@ -13,6 +14,14 @@ from fastvideo.pipelines import ForwardBatch, LoRAPipeline, build_pipeline
logger = init_logger(__name__)
def _log_cuda_device_uuid(rank: int, device: torch.device) -> None:
"""Record an NVIDIA worker UUID when external NVTX profiling is enabled."""
if not envs.FASTVIDEO_NVTX_PROFILE:
return
device_uuid = torch.cuda.get_device_properties(device).uuid
logger.info("Worker %d CUDA device UUID: GPU-%s", rank, device_uuid, local_main_process_only=False)
class Worker:
def __init__(self, fastvideo_args: FastVideoArgs, local_rank: int, rank: int, distributed_init_method: str):
@@ -61,6 +70,8 @@ class Worker:
if current_platform.is_cuda_alike():
torch.cuda.set_device(self.device)
self.init_gpu_memory = torch.cuda.mem_get_info(self.device)[0]
if current_platform.is_cuda():
_log_cuda_device_uuid(self.rank, self.device)
else:
# For MPS, we can't get memory info the same way
self.init_gpu_memory = 0