|
| 1 | +""" |
| 2 | +This script can be used to estimate training times for CPT/PPO/SFT/DPO Nova Customization training |
| 3 | +jobs on both Sagemaker HyperPod and SageMaker Training Jobs. To use the script, ensure that you have |
| 4 | +your YAML recipe that you anticipate using to start your job. |
| 5 | +
|
| 6 | +Run the script by using a command similar to: python get_training_time_estimate.py |
| 7 | +
|
| 8 | +Simply walk through the prompts to receive your training time estimate. Please note that these estimates |
| 9 | +are approximate projections and should not be interpreted as definitive training durations. Actual training |
| 10 | +times may vary significantly based on multiple factors. For more accurate estimates, please check CloudWatch |
| 11 | +logs while your training job is running. |
| 12 | +""" |
| 13 | + |
| 14 | +import os |
| 15 | +import yaml |
| 16 | + |
| 17 | +# Baselines adapted from performing regression on internal stress tests |
| 18 | +BASELINES_BY_TYPE = { |
| 19 | + "cpt": { |
| 20 | + "nova-micro": {"T0": 4.0, "N0": 100000, "P5": 8, "GBS0": 256, "MaxLength": 8192}, |
| 21 | + "nova-lite": {"T0": 4.0, "N0": 100000, "P5": 16, "GBS0": 256, "MaxLength": 8192}, |
| 22 | + "nova-pro": {"T0": 10.0, "N0": 100000, "P5": 24, "GBS0": 256, "MaxLength": 8192}, |
| 23 | + }, |
| 24 | + "ppo": { |
| 25 | + "nova-micro": {"T0": 19.6, "N0": 15000, "P5": 7, "GBS0": 160, "MaxLength": 8192}, |
| 26 | + "nova-lite": {"T0": 18.6, "N0": 20000, "P5": 7, "GBS0": 160, "MaxLength": 8192}, |
| 27 | + "nova-pro": {"T0": 52.8, "N0": 20000, "P5": 8, "GBS0": 160, "MaxLength": 8192}, |
| 28 | + }, |
| 29 | + "sft": { |
| 30 | + "nova-micro": { |
| 31 | + "full-rank": {"T0": 0.45, "N0": 5000, "P5": 2, "GBS0": 64, "MaxLength": 48000}, |
| 32 | + "lora": {"T0": 0.45, "N0": 5000, "P5": 2, "GBS0": 64, "MaxLength": 64000} |
| 33 | + }, |
| 34 | + "nova-lite": { |
| 35 | + "full-rank": {"T0": 0.5, "N0": 1500, "P5": 4, "GBS0": 64, "MaxLength": 48000}, |
| 36 | + # No Lora data exists |
| 37 | + }, |
| 38 | + "nova-pro": { |
| 39 | + "full-rank": {"T0": 0.707, "N0": 1500, "P5": 6, "GBS0": 32, "MaxLength": 48000}, |
| 40 | + "lora": {"T0": 0.75, "N0": 2000, "P5": 6, "GBS0": 32, "MaxLength": 64000} |
| 41 | + } |
| 42 | + }, |
| 43 | + "dpo": { |
| 44 | + "nova-micro": { |
| 45 | + "full-rank": {"T0": 0.46, "N0": 10500, "P5": 2, "GBS0": 256, "MaxLength": 24000}, |
| 46 | + "lora": {"T0": 0.46, "N0": 10500, "P5": 2, "GBS0": 256, "MaxLength": 24000} |
| 47 | + }, |
| 48 | + "nova-lite": { |
| 49 | + "full-rank": {"T0": 0.66, "N0": 20000, "P5": 4, "GBS0": 256, "MaxLength": 16000}, |
| 50 | + "lora": {"T0": 0.75, "N0": 20000, "P5": 4, "GBS0": 256, "MaxLength": 16000} |
| 51 | + }, |
| 52 | + "nova-pro": { |
| 53 | + "full-rank": {"T0": 1.08, "N0": 20000, "P5": 4, "GBS0": 128, "MaxLength": 16000}, |
| 54 | + "lora": {"T0": 1.5, "N0": 20000, "P5": 4, "GBS0": 128, "MaxLength": 16000} |
| 55 | + } |
| 56 | + } |
| 57 | +} |
| 58 | + |
| 59 | +# Estimate training time (in hours) based on linear scaling assumptions. |
| 60 | +def estimate_training_time_hours( |
| 61 | + model_type: str, |
| 62 | + num_samples: int, |
| 63 | + p5_instances: int, |
| 64 | + gbs: int, |
| 65 | + max_length: int, |
| 66 | + training_type: str, |
| 67 | + sub_type: str = None |
| 68 | +): |
| 69 | + # Fetch the baseline for the training type and model |
| 70 | + baseline = BASELINES_BY_TYPE[training_type][model_type] |
| 71 | + |
| 72 | + # If the baseline has subtypes (full-rank, lora, etc.), use the provided sub_type |
| 73 | + if training_type in ["sft", "dpo"]: |
| 74 | + baseline = baseline[sub_type] |
| 75 | + |
| 76 | + # Perform linear scaling calculation |
| 77 | + return ( |
| 78 | + baseline["T0"] |
| 79 | + * (num_samples / baseline["N0"]) # > samples, > time |
| 80 | + * (baseline["P5"] / p5_instances) # > P5's, < time |
| 81 | + * (baseline["GBS0"] / gbs) # > gbs, < time |
| 82 | + * (max_length / baseline["MaxLength"]) # > max_length, > time |
| 83 | + ) |
| 84 | + |
| 85 | +def retrieve_recipe(): |
| 86 | + """Prompt user for YAML file and return parsed content. Keep asking user until a valid yaml path is given""" |
| 87 | + while True: |
| 88 | + yaml_path = input("Enter path to YAML file: ").strip() |
| 89 | + if os.path.isfile(yaml_path): |
| 90 | + break |
| 91 | + print("Invalid file. Please enter a valid YAML file path.") |
| 92 | + |
| 93 | + with open(yaml_path, "r") as f: |
| 94 | + data = yaml.safe_load(f) |
| 95 | + |
| 96 | + # Sanity check: ensure required keys exist |
| 97 | + if "run" not in data or "replicas" not in data["run"]: |
| 98 | + raise ValueError("YAML missing 'run.replicas' value.") |
| 99 | + if "training_config" not in data or "global_batch_size" not in data["training_config"] or "max_length" not in data["training_config"]: |
| 100 | + raise ValueError("YAML missing 'training_config.global_batch_size' or 'training_config.max_length' values.") |
| 101 | + |
| 102 | + return data |
| 103 | + |
| 104 | +def select_from_list(prompt: str, options: list): |
| 105 | + """Display numbered list and get user selection.""" |
| 106 | + print(prompt) |
| 107 | + for i, option in enumerate(options, start=1): |
| 108 | + print(f"{i}. {option}") |
| 109 | + while True: |
| 110 | + choice = input("Enter the number of your choice: ").strip() |
| 111 | + if choice.isdigit(): |
| 112 | + idx = int(choice) - 1 |
| 113 | + if 0 <= idx < len(options): |
| 114 | + return options[idx] |
| 115 | + print(f"Invalid choice. Please enter a number between 1 and {len(options)}.") |
| 116 | + |
| 117 | +def format_days_hours_minutes(hours: float): |
| 118 | + """Convert decimal hours to 'Xd Yh Zm'""" |
| 119 | + total_minutes = int(round(hours * 60)) |
| 120 | + days = total_minutes // (24 * 60) |
| 121 | + hours_rem = (total_minutes % (24 * 60)) // 60 |
| 122 | + minutes = total_minutes % 60 |
| 123 | + |
| 124 | + def pluralize(value, unit): |
| 125 | + return f"{value} {unit}" + ("s" if value != 1 else "") |
| 126 | + |
| 127 | + parts = [] |
| 128 | + if days > 0: |
| 129 | + parts.append(pluralize(days, "day")) |
| 130 | + if hours_rem > 0: |
| 131 | + parts.append(pluralize(hours_rem, "hour")) |
| 132 | + if minutes > 0: |
| 133 | + parts.append(pluralize(minutes, "minute")) |
| 134 | + |
| 135 | + return " ".join(parts) if parts else "Less than a minute" |
| 136 | + |
| 137 | + |
| 138 | +def main(): |
| 139 | + # Present only training types that have baselines |
| 140 | + training_types = [t for t, models in BASELINES_BY_TYPE.items() if models] |
| 141 | + |
| 142 | + training_type = select_from_list("Select training type:", training_types) |
| 143 | + |
| 144 | + # Present only models available for the selected training type |
| 145 | + available_models = list(BASELINES_BY_TYPE[training_type].keys()) |
| 146 | + model_type = select_from_list("Select model type:", available_models) |
| 147 | + |
| 148 | + |
| 149 | + sub_type = None |
| 150 | + |
| 151 | + # For SFT & DPO, we must get sub_type |
| 152 | + if training_type in ["sft", "dpo"]: |
| 153 | + model_entry = BASELINES_BY_TYPE[training_type][model_type] |
| 154 | + sub_types = list(model_entry.keys()) |
| 155 | + sub_type = select_from_list("Select sub-type:", sub_types) |
| 156 | + |
| 157 | + # Prompt User for Recipe Location |
| 158 | + recipe = retrieve_recipe() |
| 159 | + |
| 160 | + # Ask for dataset size, optional |
| 161 | + dataset_input = input("Enter dataset sample size (press Enter to use default 100,000): ").strip() |
| 162 | + if dataset_input == "": |
| 163 | + num_samples = 100000 |
| 164 | + dataset_default = True |
| 165 | + else: |
| 166 | + try: |
| 167 | + num_samples = int(dataset_input) |
| 168 | + dataset_default = False |
| 169 | + except ValueError: |
| 170 | + print("Invalid input. Using default 100,000 samples.") |
| 171 | + num_samples = 100000 |
| 172 | + dataset_default = True |
| 173 | + |
| 174 | + p5_instances = recipe["run"]["replicas"] |
| 175 | + gbs = recipe["training_config"]["global_batch_size"] |
| 176 | + max_length = recipe["training_config"]["max_length"] |
| 177 | + |
| 178 | + # Estimate training time |
| 179 | + estimated_hours = estimate_training_time_hours( |
| 180 | + model_type=model_type, |
| 181 | + num_samples=num_samples, |
| 182 | + p5_instances=p5_instances, |
| 183 | + gbs=gbs, |
| 184 | + max_length=max_length, |
| 185 | + training_type=training_type, |
| 186 | + sub_type=sub_type |
| 187 | + ) |
| 188 | + |
| 189 | + # Format output |
| 190 | + formatted_time = format_days_hours_minutes(estimated_hours) |
| 191 | + |
| 192 | + # Build descriptive output string |
| 193 | + desc = f"{training_type}, {model_type}" |
| 194 | + |
| 195 | + if sub_type: |
| 196 | + desc += f", {sub_type}" |
| 197 | + |
| 198 | + desc += f", {p5_instances} P5s, {gbs} GBS, {max_length} max_length" |
| 199 | + |
| 200 | + if dataset_default: |
| 201 | + print(f"\nEstimated training time ({desc}, per 100000 samples): {formatted_time}.") |
| 202 | + else: |
| 203 | + print(f"\nEstimated training time ({desc}): {formatted_time}") |
| 204 | + |
| 205 | + print( |
| 206 | + "Please note that these estimates are approximate projections and should not be interpreted " |
| 207 | + "as definitive training durations. Please monitor CloudWatch logs for more accurate progress and estimates." |
| 208 | + ) |
| 209 | + |
| 210 | +if __name__ == "__main__": |
| 211 | + main() |
0 commit comments