diff --git a/examples/inference/gpt/gpt_dynamic_inference.py b/examples/inference/gpt/gpt_dynamic_inference.py index a3b82a69e68..84868934bc9 100644 --- a/examples/inference/gpt/gpt_dynamic_inference.py +++ b/examples/inference/gpt/gpt_dynamic_inference.py @@ -133,6 +133,10 @@ def get_inference_context( else: max_sequence_length = args.inference_max_seq_length + metrics_writer = None + if args.inference_wandb_logging_step_interval > 0: + metrics_writer = get_wandb_writer() + # Inference context. context = DynamicInferenceContext( params_dtype=args.params_dtype, @@ -161,6 +165,7 @@ def get_inference_context( use_cuda_graphs_for_non_decode_steps=not args.decode_only_cuda_graphs, use_flashinfer_fused_rope=args.use_flashinfer_fused_rope, unified_memory_level=args.inference_dynamic_batching_unified_memory_level, + metrics_writer=metrics_writer, ) return context @@ -388,6 +393,7 @@ def main(): random_seed=args.seed, track_paused_request_events=args.inference_dynamic_batching_track_paused_request_events, enable_chunked_prefill=not args.disable_chunked_prefill, + inference_logging_step_interval=args.inference_wandb_logging_step_interval, ) setup_prefix = build_dynamic_engine_setup_prefix(args, model, context, requests) diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index bdf4e9fe43a..3eece64b24f 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -4,7 +4,7 @@ import warnings from contextlib import nullcontext from enum import Enum -from typing import List, Optional, Tuple +from typing import TYPE_CHECKING, List, Optional, Tuple import torch import torch.nn.functional as F @@ -45,6 +45,17 @@ except ImportError: HAVE_FLASHINFER = False +try: + import wandb # pylint: disable=unused-import + + HAVE_WANDB = True +except ImportError: + HAVE_WANDB = False + wandb = None + +if TYPE_CHECKING: + import wandb as WandbModule + class ContextOverflowError(Exception): """Base exception for when a new request does not fit. @@ -222,6 +233,7 @@ class DynamicInferenceContext(BaseInferenceContext): levels will be included to control other tensors within the context. use_flashinfer_fused_rope (bool): If True, use flashinfer's fused rope implementation. If None, defaults to using flash-infer if available. + metrics_writer (Optional['WandbModule']): Wandb module for writing metrics. """ def __init__( @@ -247,6 +259,7 @@ def __init__( use_cuda_graphs_for_non_decode_steps: bool = True, use_flashinfer_fused_rope: bool = False, unified_memory_level: Optional[int] = 0, + metrics_writer: Optional['WandbModule'] = None, ): super().__init__(materialize_only_last_token_logits=materialize_only_last_token_logits) @@ -256,6 +269,8 @@ def __init__( block_size_tokens == 64 ), "Flash MLA requires a block size of 64. Set --inference-dynamic-batching-block-size 64 to fix this assert" + self.metrics_writer = metrics_writer + # Per partition num heads and hidden size. projection_size = kv_channels * num_attention_heads if tensor_model_parallel_size is None: @@ -1618,3 +1633,67 @@ def calculate_log_probs( # Convert each log prob tensor into a list return [lp.tolist() for lp in selected_log_probs_list] + + def get_kvcache_utilization_stats(self) -> dict: + """Compute KV cache buffer utilization stats for the current step. + + Returns a dictionary with counts and percentages for both allocated block + usage (overall buffer occupancy) and active usage (blocks referenced by + currently active requests this step). + + Return: + { + 'total_blocks': int, + 'allocated_blocks': int, + 'active_unique_blocks': int, + 'allocated_utilization': float, + 'active_utilization': float, + 'active_request_count': int, + 'paused_request_count': int, + 'gtd_block_count': int, + } + """ + # Total usable blocks exclude the reserved dummy block. + total_blocks = max(self.block_allocator.block_count_total - 1, 1) + block_count_avail = int(self.block_allocator.block_count_avail) + + # Overall allocated blocks in the buffer right now. + allocated_blocks = (self.block_allocator.block_count_total - 1) - block_count_avail + allocated_blocks = int(max(0, allocated_blocks)) + + # Active unique blocks referenced by current active requests only. + active_start = self.paused_request_count + active_end = self.total_request_count + if active_end > active_start: + active_rows = self.request_to_kv_block_ids[active_start:active_end] + # Filter valid block ids (>= 0) and count unique ids. + valid_ids = active_rows[active_rows >= 0] + if valid_ids.numel() > 0: + unique_ids = torch.unique(valid_ids) + active_unique_blocks = int(unique_ids.numel()) + else: + active_unique_blocks = 0 + else: + active_unique_blocks = 0 + + allocated_utilization = float(allocated_blocks) / float(total_blocks) + active_utilization = float(active_unique_blocks) / float(total_blocks) + + # Diagnostic helpers + num_non_gtd_blocks = max(0, block_count_avail - int(self.gtd_block_count)) + total_request_count = int(self.total_request_count) + return { + 'total_blocks': int(total_blocks), + 'allocated_blocks': int(allocated_blocks), + 'active_unique_blocks': int(active_unique_blocks), + 'allocated_utilization': allocated_utilization, + 'active_utilization': active_utilization, + 'active_request_count': int(self.get_active_request_count()), + 'paused_request_count': int(self.paused_request_count), + 'gtd_block_count': int(self.gtd_block_count), + 'block_count_avail': int(block_count_avail), + 'num_non_gtd_blocks': int(num_non_gtd_blocks), + 'active_token_count': int(self.active_token_count), + 'total_request_count': int(total_request_count), + 'max_requests': int(self.max_requests), + } diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py index b174b905e94..bcde4f9894d 100644 --- a/megatron/core/inference/engines/dynamic_engine.py +++ b/megatron/core/inference/engines/dynamic_engine.py @@ -57,6 +57,14 @@ except: HAVE_MSGPACK = False +try: + import wandb + + HAVE_WANDB = True +except ImportError: + HAVE_WANDB = False + wandb = None + def format_mem_bytes(mem_bytes): """Convert a byte count to a human-readable string in tb, gb, mb, kb, or bytes.""" @@ -89,6 +97,8 @@ class DynamicInferenceEngine(AbstractEngine): static_sampling (bool): If True, all requests are assumed to have the same sampling parameters. This avoids needing to loop through all requests and their sampling parameters every generation step, improving latency. + inference_logging_step_interval (int): The step interval at which to log + inference metrics to wandb. Defaults to 0, which means no logging. """ def __init__( @@ -101,6 +111,7 @@ def __init__( track_paused_request_events: bool = False, enable_chunked_prefill: bool = True, static_sampling: bool = False, + inference_logging_step_interval: int = 0, ): if enable_cuda_graph is not None: @@ -137,6 +148,32 @@ def __init__( self.enable_chunked_prefill = enable_chunked_prefill self.static_sampling = static_sampling + self.inference_logging_step_interval = inference_logging_step_interval + # Configure wandb to use separate step counter for inference metrics (only once) + if self.inference_logging_step_interval > 0 and self.context.metrics_writer is not None: + logging.info( + f"\033[1;93m[INFERENCE]\033[0m " + f"\033[1;95mLogging inference metrics to wandb (rank {torch.distributed.get_rank()})\033[0m" + ) + if HAVE_WANDB and self.context.metrics_writer.__name__ == "wandb": + # Make all inference/* metrics use inference_step as their x-axis + # This allows inference and training to have independent step counters + context.metrics_writer.define_metric( + "inference/*", step_metric="inference/inference_step" + ) + # Initialize inference step offset by querying existing run history + self.inference_step_offset = 0 + if wandb.run is not None: + api_run = wandb.Api().run( + f"{wandb.run.entity}/{wandb.run.project}/{wandb.run.id}" + ) + max_step = 0 + for row in api_run.scan_history(keys=["inference/inference_step"]): + val = row.get("inference/inference_step") + if isinstance(val, (int, float)) and int(val) > max_step: + max_step = int(val) + self.inference_step_offset = int(max_step) + # Initialize the asyncio loop if it has not already been initialized. # TODO: Start the engine loop here. self._loop = get_asyncio_loop() @@ -780,6 +817,41 @@ async def async_step( self.request_completion_futures[failed_request_id].set_result(failed_request) self.failed_request_ids.clear() + # Log KV cache utilization stats to W&B + if ( + self.inference_logging_step_interval > 0 + and self.step_count > 0 + and self.step_count % self.inference_logging_step_interval == 0 + and self.context.metrics_writer is not None + ): + + # Get KV cache utilization stats from dynamic context + kv_stats = self.context.get_kvcache_utilization_stats() + + # Prepare metrics dictionary with all stats + # Use 'inference/' prefix for all metrics to separate from training metrics + metrics = { + 'inference/inference_step': int(self.inference_step_offset + int(self.step_count)), + 'inference/step_time_s': float(step_time), + 'inference/waiting_queue_len': int(len(self.waiting_request_ids)), + 'inference/total_requests_dict_size': int(len(self.requests)), + } + # Add KV stats with inference/ prefix + # Convert utilization metrics from 0-1 range to 0-100 percentage range for better visualization + for key, value in kv_stats.items(): + if 'utilization' in key: + # Convert to percentage (0-100) and group under kvcache_utilization + metrics[f'inference/{key}'] = float(value * 100.0) + else: + metrics[f'inference/{key}'] = value + + if HAVE_WANDB and self.context.metrics_writer.__name__ == "wandb": + self.context.metrics_writer.log(metrics, commit=True) + else: + raise ValueError( + f"Unsupported metrics writer type: {type(self.context.metrics_writer)}" + ) + # Print context state. if verbose: context = self.context diff --git a/megatron/rl/inference/megatron.py b/megatron/rl/inference/megatron.py index 387c80d611c..5cb91d6c952 100644 --- a/megatron/rl/inference/megatron.py +++ b/megatron/rl/inference/megatron.py @@ -83,7 +83,8 @@ def get_static_inference_engine(args: Namespace, model: MegatronModule) -> Abstr ## This code is copied from tools/run_text_generation_server.py -def get_dynamic_inference_engine(args: Namespace, model: MegatronModule) -> AbstractEngine: +def get_dynamic_inference_engine(args: Namespace, model: MegatronModule, inference_logging_step_interval: int = 0, + metrics_writer = None) -> AbstractEngine: """Get the relevant backend for running inference. This function will automatically choose the TRTLLMBackend when possible, @@ -93,6 +94,8 @@ def get_dynamic_inference_engine(args: Namespace, model: MegatronModule) -> Abst Args: args (Namespace): The user arguments parsed from command line model (MegatronModule): The megatron model. + inference_logging_step_interval (int): Step interval for logging inference metrics. + metrics_writer: Metrics writer (wandb module) for logging. Returns: AbstractBackend: The chosen backend @@ -131,6 +134,7 @@ def get_dynamic_inference_engine(args: Namespace, model: MegatronModule) -> Abst mamba_d_model=args.hidden_size, mamba_d_conv=4 if args.is_hybrid_model else None, mamba_d_state=args.mamba_state_dim, + metrics_writer=metrics_writer, ) inference_wrapped_model = GPTInferenceWrapper(model, args, inference_context) @@ -148,6 +152,7 @@ def get_dynamic_inference_engine(args: Namespace, model: MegatronModule) -> Abst context=inference_context, enable_cuda_graph=args.enable_cuda_graph, random_seed=args.seed, + inference_logging_step_interval=inference_logging_step_interval, ) @@ -210,8 +215,20 @@ async def launch(cls, model: GPTModel, **kwargs): logging.WARNING, "WARNING: Tokenizer has no BOS token so prompt will not have BOS token", ) - - inference_engine: DynamicInferenceEngine = get_dynamic_inference_engine(args, model) + + # Get inference logging configuration from args + inference_logging_step_interval = args.inference_wandb_logging_step_interval + + # Get metrics writer if logging is enabled and on the logging rank + # Use the same rank convention as training (last rank logs) + metrics_writer = None + if inference_logging_step_interval > 0 and args.rank == (args.world_size - 1): + metrics_writer = get_wandb_writer() + if metrics_writer is None: + log_single_rank(logger, logging.WARNING, "WARNING: --rl-inference-logging-step-interval is set but no metrics writer " + "wandb module is available. Inference logging will be disabled.") + + inference_engine: DynamicInferenceEngine = get_dynamic_inference_engine(args, model, inference_logging_step_interval, metrics_writer) coordinator = DynamicEngineCoordinator( inference_engine, inference_max_requests=inference_engine.context.max_requests, diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index d8647607c0a..0b14140529a 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -1497,7 +1497,11 @@ def _add_inference_args(parser): help='Number of chunks along sequence dimension for MLP ' 'computation during prefill') group.add_argument('--disable-chunked-prefill', default=False, action="store_true", - help='Disable chunked prefill (chunked prefill is enabled by default).') + help='Disable chunked prefill (chunked prefill is enabled by default).') + group.add_argument('--inference-wandb-logging-step-interval', type=int, default=0, + help='Step interval for logging inference metrics to wandb. ' + 'Default to 0 to disable inference wandb logging.') + return parser diff --git a/tests/unit_tests/inference/test_wandb_logging.py b/tests/unit_tests/inference/test_wandb_logging.py new file mode 100644 index 00000000000..1512e805f9c --- /dev/null +++ b/tests/unit_tests/inference/test_wandb_logging.py @@ -0,0 +1,288 @@ +# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Unit tests for wandb logging functionality in inference.""" + +from unittest.mock import MagicMock, Mock, create_autospec, patch + +import pytest +import torch + +from megatron.core.inference.contexts.dynamic_context import DynamicInferenceContext +from megatron.core.inference.engines import DynamicInferenceEngine +from megatron.core.inference.inference_request import DynamicInferenceRequest +from megatron.core.inference.sampling_params import SamplingParams +from megatron.core.inference.text_generation_controllers.text_generation_controller import ( + TextGenerationController, +) +from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed +from tests.unit_tests.test_utilities import Utils + + +def set_rounder(value): + """Utility function to set the DynamicInferenceContext rounder.""" + DynamicInferenceContext.ROUNDER = value # For backwards compatibility + DynamicInferenceContext.TOKEN_ROUNDER = value + DynamicInferenceContext.REQUEST_ROUNDER = value + + +class TestInferenceWandbLogging: + """Test suite for wandb logging in inference.""" + + def setup_method(self): + """Set up test fixtures.""" + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1 + ) + model_parallel_cuda_manual_seed(123) + set_rounder(64) + + def teardown_method(self): + """Clean up test fixtures.""" + set_rounder(64) + Utils.destroy_model_parallel() + + def _get_dynamic_context( + self, + params_dtype=torch.float32, + num_layers=4, + kv_channels=8, + num_attention_heads=2, + max_sequence_length=512, + buffer_size_gb=0.03, + block_size_tokens=128, + buffer_guaranteed_fraction=0.1, + metrics_writer=None, + ): + """Helper to create a DynamicInferenceContext.""" + return DynamicInferenceContext( + params_dtype=params_dtype, + num_layers=num_layers, + kv_channels=kv_channels, + num_attention_heads=num_attention_heads, + max_sequence_length=max_sequence_length, + num_cuda_graphs=None, + buffer_size_gb=buffer_size_gb, + buffer_guaranteed_fraction=buffer_guaranteed_fraction, + block_size_tokens=block_size_tokens, + metrics_writer=metrics_writer, + ) + + @pytest.mark.internal + def test_get_kvcache_utilization_stats_with_requests(self): + """Test get_kvcache_utilization_stats() with empty context and then with active requests.""" + dynamic_context = self._get_dynamic_context() + + # First, test with empty context + stats = dynamic_context.get_kvcache_utilization_stats() + + # Verify all required fields are present + assert 'total_blocks' in stats + assert 'allocated_blocks' in stats + assert 'active_unique_blocks' in stats + assert 'allocated_utilization' in stats + assert 'active_utilization' in stats + assert 'active_request_count' in stats + assert 'paused_request_count' in stats + assert 'gtd_block_count' in stats + assert 'block_count_avail' in stats + assert 'num_non_gtd_blocks' in stats + assert 'active_token_count' in stats + assert 'total_request_count' in stats + assert 'max_requests' in stats + + # Verify values for empty context + assert stats['allocated_blocks'] == 0 + assert stats['active_unique_blocks'] == 0 + assert stats['allocated_utilization'] == 0.0 + assert stats['active_utilization'] == 0.0 + assert stats['active_request_count'] == 0 + assert stats['paused_request_count'] == 0 + assert stats['active_token_count'] == 0 + assert stats['total_request_count'] == 0 + + # Now add a request and verify stats update correctly + context_length = 144 + dynamic_context.add_request( + DynamicInferenceRequest( + request_id=0, + prompt_tokens=torch.arange(0, context_length, dtype=torch.long, device='cuda'), + sampling_params=SamplingParams( + num_tokens_to_generate=dynamic_context.max_tokens - context_length + ), + ) + ) + + # Initialize attention state to populate block table + dynamic_context.initialize_attention_state() + + # Get stats after adding request + stats_after = dynamic_context.get_kvcache_utilization_stats() + + # Verify that we have allocated blocks + assert stats_after['allocated_blocks'] > 0 + assert stats_after['active_unique_blocks'] > 0 + assert stats_after['allocated_utilization'] > 0.0 + assert stats_after['active_utilization'] > 0.0 + + # Verify request counts + assert stats_after['active_request_count'] == 1 + assert stats_after['total_request_count'] == 1 + assert stats_after['active_token_count'] == context_length + assert stats_after['paused_request_count'] == 0 + + # Verify that total_blocks remains constant + assert stats_after['total_blocks'] == stats['total_blocks'] + assert stats_after['total_blocks'] > 0 + + # Verify that gtd_block_count remains constant + assert stats_after['gtd_block_count'] == stats['gtd_block_count'] + + # Verify that max_requests remains constant + assert stats_after['max_requests'] == stats['max_requests'] + assert stats_after['max_requests'] > 0 + + # Verify block availability decreased after allocation + assert stats_after['block_count_avail'] < stats['block_count_avail'] + + # Verify relationship: allocated_blocks + block_count_avail + 1 (dummy) = total + assert ( + stats_after['allocated_blocks'] + stats_after['block_count_avail'] + 1 + == dynamic_context.block_allocator.block_count_total + ) + + # Verify utilization bounds [0, 1] + assert 0.0 <= stats_after['allocated_utilization'] <= 1.0 + assert 0.0 <= stats_after['active_utilization'] <= 1.0 + + # Verify relationship: active_utilization <= allocated_utilization + # (active blocks are a subset of allocated blocks) + assert stats_after['active_utilization'] <= stats_after['allocated_utilization'] + + # Verify relationship: active_unique_blocks <= allocated_blocks + assert stats_after['active_unique_blocks'] <= stats_after['allocated_blocks'] + + # Calculate expected number of blocks needed for this request + expected_blocks_needed = ( + context_length + dynamic_context.block_size_tokens - 1 + ) // dynamic_context.block_size_tokens + assert stats_after['allocated_blocks'] == expected_blocks_needed + + @pytest.mark.internal + def test_kvcache_utilization_stats_types(self): + """Test that get_kvcache_utilization_stats() returns correct types.""" + dynamic_context = self._get_dynamic_context() + stats = dynamic_context.get_kvcache_utilization_stats() + + # All integer fields + int_fields = [ + 'total_blocks', + 'allocated_blocks', + 'active_unique_blocks', + 'active_request_count', + 'paused_request_count', + 'gtd_block_count', + 'block_count_avail', + 'num_non_gtd_blocks', + 'active_token_count', + 'total_request_count', + 'max_requests', + ] + + for field in int_fields: + assert isinstance( + stats[field], int + ), f"{field} should be int but is {type(stats[field])}" + + # All float fields + float_fields = ['allocated_utilization', 'active_utilization'] + for field in float_fields: + assert isinstance( + stats[field], float + ), f"{field} should be float but is {type(stats[field])}" + + @pytest.mark.internal + @patch('megatron.core.inference.engines.dynamic_engine.HAVE_WANDB', True) + def test_engine_logging_step_interval_zero(self): + """Test that no logging occurs when inference_logging_step_interval is 0.""" + mock_wandb = Mock() + mock_wandb.__name__ = "wandb" + mock_wandb.log = Mock() + + dynamic_context = self._get_dynamic_context(metrics_writer=mock_wandb) + + # Create mock controller with proper spec to pass isinstance checks + mock_controller = create_autospec(TextGenerationController, instance=True) + # Set up nested mock structure + mock_controller.inference_wrapped_model = Mock() + mock_controller.inference_wrapped_model.model = Mock() + mock_controller.inference_wrapped_model.model.config = Mock() + mock_controller.inference_wrapped_model.model.config.cuda_graph_impl = "none" + + engine = DynamicInferenceEngine( + controller=mock_controller, + context=dynamic_context, + random_seed=123, + inference_logging_step_interval=0, # Disabled + ) + + # Verify log was never called + mock_wandb.log.assert_not_called() + + @pytest.mark.internal + def test_paused_requests_in_stats(self): + """Test that paused requests are correctly reflected in stats.""" + set_rounder(1) + dynamic_context = DynamicInferenceContext( + params_dtype=torch.float32, + num_layers=2, + kv_channels=64, + num_attention_heads=8, + max_sequence_length=128, + num_cuda_graphs=None, + buffer_size_gb=0.01, # Small buffer to force pausing + buffer_guaranteed_fraction=0.1, + block_size_tokens=32, + ) + + # Add multiple requests to potentially trigger pausing + for i in range(5): + dynamic_context.add_request( + DynamicInferenceRequest( + request_id=i, + prompt_tokens=torch.zeros(10, device='cuda'), + sampling_params=SamplingParams(num_tokens_to_generate=10), + ) + ) + + if dynamic_context.total_request_count > 0: + dynamic_context.initialize_attention_state() + stats = dynamic_context.get_kvcache_utilization_stats() + + # Verify paused request count is included + assert 'paused_request_count' in stats + assert stats['paused_request_count'] >= 0 + + @pytest.mark.internal + def test_metrics_writer_none_handling(self): + """Test that engine handles None metrics_writer gracefully.""" + dynamic_context = self._get_dynamic_context(metrics_writer=None) + + # Create mock controller with proper spec to pass isinstance checks + mock_controller = create_autospec(TextGenerationController, instance=True) + # Set up nested mock structure + mock_controller.inference_wrapped_model = Mock() + mock_controller.inference_wrapped_model.model = Mock() + mock_controller.inference_wrapped_model.model.config = Mock() + mock_controller.inference_wrapped_model.model.config.cuda_graph_impl = "none" + + # Should not raise error even with logging interval set + engine = DynamicInferenceEngine( + controller=mock_controller, + context=dynamic_context, + random_seed=123, + inference_logging_step_interval=10, + ) + + # Verify engine was created successfully + assert engine.inference_logging_step_interval == 10 + assert engine.context.metrics_writer is None