From 0aaa64fbd8176fe85a10dc237a88be0c0ec2c7d1 Mon Sep 17 00:00:00 2001 From: Anna Shors Date: Fri, 31 Oct 2025 14:23:37 -0700 Subject: [PATCH 1/2] fix: Fixes to make Megatron backend match dtensor (#1389) Signed-off-by: ashors1 Signed-off-by: Anna Shors --- examples/configs/distillation_math.yaml | 1 - .../configs/distillation_math_megatron.yaml | 1 - examples/configs/dpo.yaml | 1 - examples/configs/grpo_math_1B.yaml | 1 - examples/configs/grpo_math_1B_megatron.yaml | 1 - .../llm/sft-qwen2.5-math7b-2n8g-megatron.yaml | 53 +++++ examples/configs/rm.yaml | 1 - examples/configs/sft.yaml | 1 - .../sft_openmathinstruct2_megatron.yaml | 1 - examples/configs/vlm_grpo_3B.yaml | 1 - examples/configs/vlm_grpo_3B_megatron.yaml | 1 - nemo_rl/models/policy/__init__.py | 3 +- .../models/policy/megatron_policy_worker.py | 66 +++++- .../llm/sft-qwen2.5-math7b-2n8g-megatron.sh | 43 ++++ tests/test_suites/nightly.txt | 2 + .../models/generation/test_vllm_generation.py | 1 - .../models/policy/test_megatron_worker.py | 199 +++++++++++++++++- tools/refit_verifier.py | 1 - 18 files changed, 360 insertions(+), 18 deletions(-) create mode 100644 examples/configs/recipes/llm/sft-qwen2.5-math7b-2n8g-megatron.yaml create mode 100755 tests/test_suites/llm/sft-qwen2.5-math7b-2n8g-megatron.sh diff --git a/examples/configs/distillation_math.yaml b/examples/configs/distillation_math.yaml index 92ae09d8ee..9ece52647b 100644 --- a/examples/configs/distillation_math.yaml +++ b/examples/configs/distillation_math.yaml @@ -145,7 +145,6 @@ policy: &POLICY_BASE grad_reduce_in_fp32: false overlap_grad_reduce: true overlap_param_gather: true - average_in_collective: true use_custom_fsdp: false data_parallel_sharding_strategy: "optim_grads_params" diff --git a/examples/configs/distillation_math_megatron.yaml b/examples/configs/distillation_math_megatron.yaml index 3df59eba84..b005f13c29 100644 --- a/examples/configs/distillation_math_megatron.yaml +++ b/examples/configs/distillation_math_megatron.yaml @@ -99,7 +99,6 @@ policy: &POLICY_BASE grad_reduce_in_fp32: false overlap_grad_reduce: true overlap_param_gather: true - average_in_collective: true use_custom_fsdp: false data_parallel_sharding_strategy: "optim_grads_params" diff --git a/examples/configs/dpo.yaml b/examples/configs/dpo.yaml index 72823364cf..2d153bc88b 100755 --- a/examples/configs/dpo.yaml +++ b/examples/configs/dpo.yaml @@ -153,7 +153,6 @@ policy: grad_reduce_in_fp32: false overlap_grad_reduce: true overlap_param_gather: true - average_in_collective: true data_parallel_sharding_strategy: "optim_grads_params" data: diff --git a/examples/configs/grpo_math_1B.yaml b/examples/configs/grpo_math_1B.yaml index fbb10ed3a6..873cb0629b 100644 --- a/examples/configs/grpo_math_1B.yaml +++ b/examples/configs/grpo_math_1B.yaml @@ -129,7 +129,6 @@ policy: grad_reduce_in_fp32: false overlap_grad_reduce: true overlap_param_gather: true - average_in_collective: true use_custom_fsdp: false data_parallel_sharding_strategy: "optim_grads_params" diff --git a/examples/configs/grpo_math_1B_megatron.yaml b/examples/configs/grpo_math_1B_megatron.yaml index 01f33f63d6..8d28d878b7 100644 --- a/examples/configs/grpo_math_1B_megatron.yaml +++ b/examples/configs/grpo_math_1B_megatron.yaml @@ -133,7 +133,6 @@ policy: grad_reduce_in_fp32: false overlap_grad_reduce: true overlap_param_gather: true - average_in_collective: true use_custom_fsdp: false data_parallel_sharding_strategy: "optim_grads_params" diff --git a/examples/configs/recipes/llm/sft-qwen2.5-math7b-2n8g-megatron.yaml b/examples/configs/recipes/llm/sft-qwen2.5-math7b-2n8g-megatron.yaml new file mode 100644 index 0000000000..151319df3a --- /dev/null +++ b/examples/configs/recipes/llm/sft-qwen2.5-math7b-2n8g-megatron.yaml @@ -0,0 +1,53 @@ +defaults: ../../sft.yaml +sft: + max_num_steps: 80 +checkpointing: + enabled: false +policy: + model_name: Qwen/Qwen2.5-Math-7B + train_global_batch_size: 512 + max_total_sequence_length: 16384 + dtensor_cfg: + enabled: false + megatron_cfg: + enabled: true + tensor_model_parallel_size: 4 + context_parallel_size: 2 + sequence_parallel: true + freeze_moe_router: true + moe_router_dtype: fp64 + moe_router_bias_update_rate: 0.0 + moe_permute_fusion: true + optimizer: + lr: 1.0e-06 + min_lr: 1.0e-06 + bf16: true + adam_beta2: 0.999 + adam_eps: 1.0e-08 + use_distributed_optimizer: false + use_precision_aware_optimizer: false + scheduler: + lr_decay_iters: null + lr_warmup_iters: 10 + lr_warmup_init: 1.0e-11 + sequence_packing: + enabled: true + make_sequence_length_divisible_by: 32 +data: + dataset_name: openmathinstruct2 + prompt_file: examples/prompts/math.txt + split: train_1M + add_generation_prompt: true + output_key: generated_solution + num_workers: 8 +logger: + wandb: + project: nemo-rl + name: sft-qwen2.5-math-7b-megatron + tensorboard: + log_dir: tb_logs-sft-qwen2.5-math-7b-megatron + mlflow: + run_name: sft-qwen2.5-math-7b-megatron +cluster: + gpus_per_node: 8 + num_nodes: 2 diff --git a/examples/configs/rm.yaml b/examples/configs/rm.yaml index 22e22ac5dd..b05b568318 100644 --- a/examples/configs/rm.yaml +++ b/examples/configs/rm.yaml @@ -123,7 +123,6 @@ policy: grad_reduce_in_fp32: false overlap_grad_reduce: true overlap_param_gather: false - average_in_collective: true data_parallel_sharding_strategy: "optim_grads_params" diff --git a/examples/configs/sft.yaml b/examples/configs/sft.yaml index b742846d2d..3e465de49b 100644 --- a/examples/configs/sft.yaml +++ b/examples/configs/sft.yaml @@ -132,7 +132,6 @@ policy: grad_reduce_in_fp32: false overlap_grad_reduce: true overlap_param_gather: true - average_in_collective: true data_parallel_sharding_strategy: "optim_grads_params" use_custom_fsdp: false diff --git a/examples/configs/sft_openmathinstruct2_megatron.yaml b/examples/configs/sft_openmathinstruct2_megatron.yaml index 696d025976..f37e8f478b 100644 --- a/examples/configs/sft_openmathinstruct2_megatron.yaml +++ b/examples/configs/sft_openmathinstruct2_megatron.yaml @@ -35,7 +35,6 @@ policy: activation_checkpointing: false context_parallel_size: 1 distributed_data_parallel_config: - average_in_collective: true data_parallel_sharding_strategy: optim_grads_params grad_reduce_in_fp32: true overlap_grad_reduce: true diff --git a/examples/configs/vlm_grpo_3B.yaml b/examples/configs/vlm_grpo_3B.yaml index 67e1899c65..ae3ec84a18 100644 --- a/examples/configs/vlm_grpo_3B.yaml +++ b/examples/configs/vlm_grpo_3B.yaml @@ -118,7 +118,6 @@ policy: grad_reduce_in_fp32: false overlap_grad_reduce: true overlap_param_gather: true - average_in_collective: true use_custom_fsdp: false data_parallel_sharding_strategy: "optim_grads_params" diff --git a/examples/configs/vlm_grpo_3B_megatron.yaml b/examples/configs/vlm_grpo_3B_megatron.yaml index 4f53358fb3..dca1e485bf 100644 --- a/examples/configs/vlm_grpo_3B_megatron.yaml +++ b/examples/configs/vlm_grpo_3B_megatron.yaml @@ -147,7 +147,6 @@ policy: grad_reduce_in_fp32: false overlap_grad_reduce: false overlap_param_gather: true - average_in_collective: true use_custom_fsdp: false data_parallel_sharding_strategy: optim_grads_params data: diff --git a/nemo_rl/models/policy/__init__.py b/nemo_rl/models/policy/__init__.py index 994a1d41ae..dd43501e81 100644 --- a/nemo_rl/models/policy/__init__.py +++ b/nemo_rl/models/policy/__init__.py @@ -73,7 +73,7 @@ class MegatronSchedulerConfig(TypedDict): end_weight_decay: float weight_decay_incr_style: str lr_decay_style: str - lr_decay_iters: int + lr_decay_iters: int | None lr_warmup_iters: int lr_warmup_init: float @@ -82,7 +82,6 @@ class MegatronDDPConfig(TypedDict): grad_reduce_in_fp32: bool overlap_grad_reduce: bool overlap_param_gather: bool - average_in_collective: bool use_custom_fsdp: bool data_parallel_sharding_strategy: str diff --git a/nemo_rl/models/policy/megatron_policy_worker.py b/nemo_rl/models/policy/megatron_policy_worker.py index f7fb8b2ed8..59760ff067 100644 --- a/nemo_rl/models/policy/megatron_policy_worker.py +++ b/nemo_rl/models/policy/megatron_policy_worker.py @@ -657,6 +657,22 @@ def __init__( "https://github.com/NVIDIA-NeMo/RL/blob/bccbc377705a81a1f4b3c31ad9767bcc15f735a8/nemo_rl/algorithms/sft.py#L175-L179." ) + ## These settings are required for correct gradient computations in mcore + ## when calculate_per_token_loss is True, there is no scaling of the gradient in mcore, + ## so we handle the scaling in nemo-rl. + ## perform_initialization = True is a workaround to ensure the correct tensor parallel attributes are set + ## on the TP-sharded parameters. + model_cfg.calculate_per_token_loss = True + model_cfg.perform_initialization = True + + assert ( + "aux_loss" not in model_cfg.moe_router_load_balancing_type + or model_cfg.moe_aux_loss_coeff == 0 + ), ( + "MoE aux loss is currently not supported due to a known bug in Megatron-LM. " + "See https://github.com/NVIDIA/Megatron-LM/issues/1984 for more details." + ) + self.megatron_cfg = ConfigContainer( model=model_cfg, checkpoint=checkpoint_config, @@ -682,9 +698,9 @@ def __init__( overlap_param_gather=self.cfg["megatron_cfg"][ "distributed_data_parallel_config" ]["overlap_param_gather"], - average_in_collective=self.cfg["megatron_cfg"][ - "distributed_data_parallel_config" - ]["average_in_collective"], + # we need to set average_in_collective=False with calculate_per_token_loss=True. + # otherwise, mcore throws an assertion error. + average_in_collective=False, use_distributed_optimizer=self.cfg["megatron_cfg"]["optimizer"][ "use_distributed_optimizer" ], @@ -2224,3 +2240,47 @@ def report_node_ip_and_gpu_id(self) -> list[tuple[str, int]]: ip = ray._private.services.get_node_ip_address() gpu_id = ray.get_gpu_ids()[0] return (ip, gpu_id) + + def check_tensor_parallel_attributes(self) -> dict[str, Any]: + """Check tensor parallel attributes on model parameters. + + Returns: + Dictionary containing information about tensor parallel parameters: + - tp_params: List of parameter names that have tensor_model_parallel=True + - non_tp_params: List of parameter names that have tensor_model_parallel=False + - total_params: Total number of parameters checked + - tp_size: Tensor parallel size from config + """ + tp_params = [] + non_tp_params = [] + total_params = 0 + + for name, param in self.model.named_parameters(): + total_params += 1 + tensor_model_parallel = getattr(param, "tensor_model_parallel", False) + + if tensor_model_parallel: + tp_params.append( + { + "name": name, + "tensor_model_parallel": tensor_model_parallel, + "partition_dim": getattr(param, "partition_dim", None), + "partition_stride": getattr(param, "partition_stride", None), + "shape": list(param.shape), + } + ) + else: + non_tp_params.append( + { + "name": name, + "tensor_model_parallel": tensor_model_parallel, + "shape": list(param.shape), + } + ) + + return { + "tp_params": tp_params, + "non_tp_params": non_tp_params, + "total_params": total_params, + "tp_size": self.megatron_cfg.model.tensor_model_parallel_size, + } diff --git a/tests/test_suites/llm/sft-qwen2.5-math7b-2n8g-megatron.sh b/tests/test_suites/llm/sft-qwen2.5-math7b-2n8g-megatron.sh new file mode 100755 index 0000000000..897f4fbb60 --- /dev/null +++ b/tests/test_suites/llm/sft-qwen2.5-math7b-2n8g-megatron.sh @@ -0,0 +1,43 @@ +#!/bin/bash +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) +source $SCRIPT_DIR/common.env + +# TODO: this config can crash on OOM +# https://github.com/NVIDIA-NeMo/RL/issues/263 + +# ===== BEGIN CONFIG ===== +NUM_NODES=2 +STEPS_PER_RUN=80 # step_time ~ 29sec +MAX_STEPS=80 +NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN )) # Round up +NUM_MINUTES=30 +# ===== END CONFIG ===== + +exit_if_max_steps_reached + +# Run the experiment +cd $PROJECT_ROOT +uv run examples/run_sft.py \ + --config $CONFIG_PATH \ + sft.max_num_steps=$MAX_STEPS \ + logger.log_dir=$LOG_DIR \ + logger.wandb_enabled=True \ + logger.wandb.project=nemo-rl \ + logger.wandb.name=$EXP_NAME \ + logger.monitor_gpus=True \ + logger.tensorboard_enabled=True \ + checkpointing.enabled=True \ + checkpointing.checkpoint_dir=$CKPT_DIR \ + ~policy.tokenizer.chat_template \ + $@ \ + 2>&1 | tee $RUN_LOG + +# Convert tensorboard logs to json +uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS + +# Only run metrics if the target step is reached +if [[ $(jq 'to_entries | .[] | select(.key == "train/loss") | .value | keys | map(tonumber) | max' $JSON_METRICS) -ge $MAX_STEPS ]]; then + uv run tests/check_metrics.py $JSON_METRICS \ + 'data["train/loss"]["80"] < 0.301' \ + 'data["validation/val_loss"]["80"] < 0.304' +fi diff --git a/tests/test_suites/nightly.txt b/tests/test_suites/nightly.txt index 4a09c6b92d..8e378024f9 100644 --- a/tests/test_suites/nightly.txt +++ b/tests/test_suites/nightly.txt @@ -66,6 +66,8 @@ tests/test_suites/llm/sft-qwen2.5-32b-4n8g-fsdp2tp8sp-actckpt.v3.sh tests/test_suites/llm/sft-llama3.1-8b-1n8g-megatron.sh # sequence packing tests/test_suites/llm/sft-llama3.1-8b-1n8g-megatron-seqpack.sh +# validate TP/DP +tests/test_suites/llm/sft-qwen2.5-math7b-2n8g-megatron.sh ####### # DPO # diff --git a/tests/unit/models/generation/test_vllm_generation.py b/tests/unit/models/generation/test_vllm_generation.py index f1846477f9..2624c80d5b 100644 --- a/tests/unit/models/generation/test_vllm_generation.py +++ b/tests/unit/models/generation/test_vllm_generation.py @@ -208,7 +208,6 @@ def get_basic_megatron_test_config( "grad_reduce_in_fp32": False, "overlap_grad_reduce": True, "overlap_param_gather": False, - "average_in_collective": True, "data_parallel_sharding_strategy": "optim_grads_params", }, }, diff --git a/tests/unit/models/policy/test_megatron_worker.py b/tests/unit/models/policy/test_megatron_worker.py index e47979788a..bd8bb170b2 100644 --- a/tests/unit/models/policy/test_megatron_worker.py +++ b/tests/unit/models/policy/test_megatron_worker.py @@ -16,7 +16,9 @@ import time from typing import Optional +import numpy as np import pytest +import ray import torch from nemo_rl.algorithms.interfaces import LossFunction @@ -148,7 +150,6 @@ def create_megatron_test_config( "grad_reduce_in_fp32": False, "overlap_grad_reduce": True, "overlap_param_gather": False, - "average_in_collective": True, "data_parallel_sharding_strategy": "optim_grads_params", }, "fp8_cfg": { @@ -2266,6 +2267,202 @@ def test_megatron_context_parallel_training_agreement(tiny_llama_model_path): ) +@pytest.mark.hf_gated +@pytest.mark.timeout(300) +def test_megatron_gradient_norm_consistency_across_parallelism(tiny_llama_model_path): + """Test that gradient norms are consistent across different TP and DP configurations. + + This test validates that the same model produces identical gradient norms + regardless of tensor parallelism (TP) and data parallelism (DP) settings. + """ + batch_size = 8 + seq_len = 64 + vocab_size = 32000 + + # Create reproducible test data + torch.manual_seed(42) + input_ids = torch.randint(0, vocab_size, (batch_size, seq_len)) + attention_mask = torch.ones(batch_size, seq_len) + input_lengths = attention_mask.sum(dim=1).to(torch.int32) + labels = torch.randint(0, vocab_size, (batch_size, seq_len)) + + data = BatchedDataDict( + { + "input_ids": input_ids, + "input_lengths": input_lengths, + "attention_mask": attention_mask, + "labels": labels, + "sample_mask": torch.ones(batch_size), + "token_mask": torch.ones_like(input_ids), + } + ) + + # Test configurations: (num_gpus, tp, pp, description) + test_configs = [ + (1, 1, 1, "DP1TP1"), + (2, 1, 1, "DP2"), # Data parallel with 2 GPUs + (2, 2, 1, "TP2"), # Tensor parallel with 2 GPUs + ] + + grad_norms = {} + losses = {} + + for num_gpus, tp, pp, desc in test_configs: + print( + f"\n=== Testing {desc} configuration (GPUs={num_gpus}, TP={tp}, PP={pp}) ===" + ) + + cluster = RayVirtualCluster( + name=f"test-grad-norm-{desc.lower()}", + bundle_ct_per_node_list=[num_gpus], + use_gpus=True, + num_gpus_per_node=num_gpus, + max_colocated_worker_groups=1, + ) + + config = create_megatron_test_config( + model_name=tiny_llama_model_path, + tp=tp, + pp=pp, + precision="float32", # Use float32 for more stable gradient comparisons + ) + + tokenizer = get_tokenizer(config["tokenizer"]) + config["generation"] = configure_generation_config( + config["generation"], tokenizer + ) + + policy = Policy( + cluster=cluster, + config=config, + tokenizer=tokenizer, + init_reference_model=False, + ) + + # Use SimpleLoss for consistent comparison + loss_fn = NLLLoss() + + try: + # Prepare for training + policy.prepare_for_training() + + # Perform one forward/backward step + print(f"Performing forward/backward pass for {desc}...") + results = policy.train(data, loss_fn) + + # Extract metrics + loss_tensor = results["loss"] + grad_norm = results["grad_norm"] + + # Verify loss is valid + assert not torch.isnan(loss_tensor).any(), ( + f"Loss should not be NaN for {desc}" + ) + assert not torch.isinf(loss_tensor).any(), ( + f"Loss should not be Inf for {desc}" + ) + + # Store results for comparison + grad_norms[desc] = grad_norm + losses[desc] = loss_tensor.cpu().numpy() + + print(f"{desc} - Loss: {loss_tensor}") + print(f"{desc} - Grad norm: {grad_norm}") + + # Check tensor parallel attributes on model parameters + print(f"Checking tensor parallel attributes for {desc}...") + tp_check_futures = policy.worker_group.run_all_workers_single_data( + "check_tensor_parallel_attributes" + ) + tp_check_results = [ray.get(future) for future in tp_check_futures] + + # Analyze the first worker's results (all workers should have the same structure) + tp_info = tp_check_results[0] + + print(f"{desc} - TP size: {tp_info['tp_size']}") + print(f"{desc} - Total params: {tp_info['total_params']}") + print(f"{desc} - TP params: {len(tp_info['tp_params'])}") + print(f"{desc} - Non-TP params: {len(tp_info['non_tp_params'])}") + + # Validate tensor parallel attributes + expected_tp_size = tp + assert tp_info["tp_size"] == expected_tp_size, ( + f"Expected TP size {expected_tp_size}, got {tp_info['tp_size']}" + ) + + if tp > 1: + tp_sharded_names = [item["name"] for item in tp_info["tp_params"]] + # When tensor parallelism is enabled, we should have some TP parameters + assert "module.embedding.word_embeddings.weight" in tp_sharded_names, ( + f"Expected module.embedding.word_embeddings.weight to be TP-sharded when TP={tp}" + ) + + finally: + policy.shutdown() + cluster.shutdown() + + # Compare gradient norms across configurations + print("\n=== Comparing gradient norms across configurations ===") + + # Get reference values from DP2 configuration + # NOTE: even if TP2 config passes these tests, it doesn't necessarily imply + # there are no bugs. That's why we also check that TP attributes are set correctly above + reference_config = "DP1TP1" + reference_grad_norm = grad_norms[reference_config] + reference_loss = losses[reference_config] + + for config_name, grad_norm in grad_norms.items(): + if config_name == reference_config: + continue + + if not isinstance(grad_norm, list): + grad_norm = [grad_norm] + + print(f"\nComparing {config_name} with {reference_config}:") + print(f" {reference_config} grad norm: {reference_grad_norm}") + print(f" {config_name} grad norm: {grad_norm}") + + # Compare gradient norms + if not isinstance(grad_norm, list): + grad_norm = [grad_norm] + reference_grad_norm = [reference_grad_norm] + if isinstance(grad_norm, list) and isinstance(reference_grad_norm, list): + # Handle case where grad_norm is a list (multiple microbatches) + assert len(grad_norm) == len(reference_grad_norm), ( + f"Number of gradient norm values should match: {len(grad_norm)} vs {len(reference_grad_norm)}" + ) + + for i, (gn, ref_gn) in enumerate(zip(grad_norm, reference_grad_norm)): + grad_diff = abs(gn - ref_gn) + relative_diff = grad_diff / (ref_gn + 1e-8) + print( + f" Microbatch {i}: {ref_gn} vs {gn}, diff={grad_diff.item():.6f}, rel_diff={relative_diff.item():.6f}" + ) + + # Allow small differences due to floating point precision and parallelization + assert relative_diff < 0.01 or grad_diff < 1e-6, ( + f"Gradient norm difference too large for microbatch {i}: " + f"{ref_gn} vs {gn} (diff={grad_diff.item():.6f}, rel_diff={relative_diff.item():.6f})" + ) + + # Compare losses (should also be identical for same computation) + loss_diff = np.max(np.abs(reference_loss - losses[config_name])) + relative_loss_diff = loss_diff / (np.mean(np.abs(reference_loss)) + 1e-8) + print( + f" Loss diff: {loss_diff:.6f}, relative loss diff: {relative_loss_diff:.6f}" + ) + + # Allow small differences in loss as well + assert relative_loss_diff < 0.01 or loss_diff < 1e-6, ( + f"Loss difference too large: " + f"max diff={loss_diff:.6f}, rel_diff={relative_loss_diff:.6f}" + ) + + print( + "\n✓ SUCCESS: Gradient norms are consistent across all parallelization configurations!" + ) + + @pytest.mark.hf_gated @pytest.mark.timeout(300) def test_megatron_policy_flops_range_check(tiny_llama_model_path): diff --git a/tools/refit_verifier.py b/tools/refit_verifier.py index 475bfd1215..ee21370797 100644 --- a/tools/refit_verifier.py +++ b/tools/refit_verifier.py @@ -249,7 +249,6 @@ def setup_configs(args, tokenizer): "grad_reduce_in_fp32": False, "overlap_grad_reduce": False, "overlap_param_gather": False, - "average_in_collective": False, "use_custom_fsdp": False, "data_parallel_sharding_strategy": "optim_grads_params", }, From fcf7285e180d74a5aa49bf578ee81189bd17ad01 Mon Sep 17 00:00:00 2001 From: Terry Kong Date: Tue, 4 Nov 2025 17:40:35 +0000 Subject: [PATCH 2/2] 1040 -> 1050 hours Signed-off-by: Terry Kong --- tests/unit/test_recipes_and_test_suites.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/unit/test_recipes_and_test_suites.py b/tests/unit/test_recipes_and_test_suites.py index 5fc984e246..d8201730b9 100644 --- a/tests/unit/test_recipes_and_test_suites.py +++ b/tests/unit/test_recipes_and_test_suites.py @@ -187,8 +187,8 @@ def test_nightly_compute_stays_below_1040_hours(nightly_test_suite, tracker): f"Last line of output was not as expected: '{last_line}'" ) total_gpu_hours = float(last_line.split(":")[-1].strip()) - assert total_gpu_hours <= 1040, ( - f"Total GPU hours exceeded 1040: {last_line}. We should revisit the test suites to reduce the total GPU hours." + assert total_gpu_hours <= 1050, ( + f"Total GPU hours exceeded 1050: {last_line}. We should revisit the test suites to reduce the total GPU hours." ) tracker.track("total_nightly_gpu_hours", total_gpu_hours)