Scaling Plots & Analysis as an Executor Step

docs/references/default-steps.md

@@ -27,12 +27,6 @@ In general, you should reach for the default steps before writing your own.
 
 ::: experiments.defaults.simulated_epoching_train
 
-## Scaling Law Prediction
-
-::: marin.scaling_laws.create_ladder_suite.scaling_law_suite
-
-::: experiments.defaults.default_scaling_law_pred
-
 ## Evaluation
 
 ::: experiments.evals.evals.default_eval

experiments/defaults.py

@@ -46,7 +46,6 @@
     CORE_TASKS,
     MMLU_TASKS,
     convert_to_levanter_task_config,
-    convert_to_task_metrics,
 )
 from experiments.llama import compute_num_parameters, llama_8b
 from experiments.paloma import paloma_tokenized
@@ -59,7 +58,6 @@
     InputName,
     VersionedValue,
     ensure_versioned,
-    get_executor_step,
     this_output_path,
     unwrap_versioned_value,
 )
@@ -72,7 +70,6 @@
     tokenize,
 )
 from marin.processing.tokenize.tokenize import HfTokenizeConfig, TokenizeConfigBase
-from marin.scaling_laws.scaling_laws import ScalingLawConfig, run_scaling_law_analysis
 from marin.training.training import (
     TrainLmOnPodConfig,
     run_levanter_train_lm,
@@ -637,41 +634,3 @@ def _get_tokenizer_for_train(tokenized: InputName | ExecutorStep | LMMixtureData
             raise ValueError(f"Could not determine tokenizer from {tokenized}")
 
     return tokenizer
-
-
-def default_scaling_law_pred(
-    ladder_runs: Sequence[ExecutorStep | InputName | str],
-    pred_run: ExecutorStep | InputName | str | None = None,
-    task_losses: Sequence[str] = ("eval/paloma/c4_en/bpb",),
-    task_accuracies: Sequence[str] | Sequence[EvalTaskConfig] | None = None,
-):
-    """
-    Given a suite of small models, predict the performance on a number of (N, D) values.
-    """
-    # get the executor steps or run IDs for the ladder runs and the pred run
-    ladder_steps_or_ids = [get_executor_step(run) if not isinstance(run, str) else run for run in ladder_runs]
-
-    pred_run_or_id = None
-    if pred_run:
-        pred_run_or_id = get_executor_step(pred_run) if not isinstance(pred_run, str) else pred_run
-
-    # convert the task accuracies to strings if they are `EvalTaskConfig`s
-    if task_accuracies is not None:
-        task_accuracies = convert_to_task_metrics(task_accuracies, metric="acc")
-
-    if pred_run_or_id:
-        name = pred_run_or_id if isinstance(pred_run_or_id, str) else pred_run_or_id.name
-    else:
-        name = "projection"
-
-    return ExecutorStep(
-        name=f"""scaling_laws/{name}""",
-        fn=run_scaling_law_analysis,
-        config=ScalingLawConfig(
-            name=name,
-            ladder_model_steps=ladder_steps_or_ids,
-            pred_model_step=pred_run_or_id,
-            task_losses=task_losses,
-            task_accuracies=task_accuracies,
-        ),
-    )

experiments/exp1600_perpcorr.py

@@ -24,7 +24,7 @@
 
 from experiments.evals.evals import evaluate_levanter_lm_evaluation_harness
 from experiments.evals.task_configs import EvalTaskConfig
-from experiments.isoflop_sweep import generate_isoflop_sweep
+from experiments.isoflop_sweep import MARIN_2025_RECIPE, create_isoflop_sweep_steps
 from experiments.llama import llama3_tokenizer
 from experiments.models import ModelConfig as HFModelConfig, download_model_step
 from experiments.paloma import paloma_tokenized
@@ -56,22 +56,22 @@
 @lru_cache(maxsize=1)
 def build_steps():
     steps = []
-    isoflop_steps, isoflop_metadatas = generate_isoflop_sweep(
+    isoflop_steps, isoflop_candidates = create_isoflop_sweep_steps(
         nemotron_mix,
         experiment_name="nemo-wider-depth-adapt",
+        recipe=MARIN_2025_RECIPE,
     )
-    for isoflop_step, isoflop_metadata in zip(isoflop_steps, isoflop_metadatas, strict=False):
+    for isoflop_step, candidate in zip(isoflop_steps, isoflop_candidates, strict=False):
         experiment_name = isoflop_step.name.split("/")[-1]
         paloma_tokenized_dict = paloma_tokenized(tokenizer=llama3_tokenizer)
         uncheatable_eval_tokenized_dict = uncheatable_eval_tokenized(tokenizer=llama3_tokenizer)
         eval_data = mixture_for_evaluation(paloma_tokenized_dict | uncheatable_eval_tokenized_dict)
-        budget, hidden_size, num_layers, batch_size, train_steps = isoflop_metadata
         wandb_tags = [
-            f"FLOPs={budget:.1e}",
-            f"d={hidden_size}",
-            f"L={num_layers}",
-            f"B={batch_size}",
-            f"steps={train_steps}",
+            f"FLOPs={candidate.flops_budget:.1e}",
+            f"d={candidate.hidden_size}",
+            f"L={candidate.num_layers}",
+            f"B={candidate.batch_size}",
+            f"steps={candidate.train_steps}",
         ]
         model_config = isoflop_step.config.train_config.model
         checkpoint_path = output_path_of(isoflop_step)

experiments/exp1603_subgroup_evals.py

@@ -24,6 +24,10 @@
 from experiments.models import ModelConfig, download_model_step
 from marin.execution.executor import executor_main, output_path_of, versioned
 from marin.evaluation.log_probs import default_lm_log_probs
+from marin.processing.tokenize import get_vocab_size_for_tokenizer
+
+# Vocab size for building model configs
+VOCAB_SIZE = get_vocab_size_for_tokenizer("stanford-crfm/marin-tokenizer")
 
 # This is painfully slow to run in dry run mode
 # nodryrun
@@ -40,8 +44,10 @@ def create_eval_steps() -> list:
 
     steps = []
     dist_eval = distributional_eval_sets(llama3_tokenizer)
-    for model, metadata in list(zip(*MARIN_SCALING_SUITES["nemotron"], strict=False)):
-        name = f"marin-nemo-{metadata[0]}C-{metadata[-3] * metadata[-2] * 4096}T-{metadata[1]}W-{metadata[2]}D"
+    for model, candidate in list(zip(*MARIN_SCALING_SUITES["nemotron"], strict=False)):
+        total_tokens = int(candidate.tokens)
+        params = candidate.model_config.total_trainable_params(VOCAB_SIZE)
+        name = f"marin-nemo-{candidate.flops_budget:.0e}C-{total_tokens}T-N{params:.0e}"
 
         step = evaluate_levanter_lm_evaluation_harness(
             model_name=name,
@@ -51,9 +57,10 @@ def create_eval_steps() -> list:
         )
         steps.append(step)
 
+        model_config = candidate.model_config
         logprobs_step = default_lm_log_probs(
             output_path_of(model).cd("checkpoints"),
-            metadata[-1],
+            model_config,
             dist_eval,
             resource_config=ResourceConfig.with_tpu("v5p-8"),
             checkpoint_is_hf=False,
@@ -62,8 +69,10 @@ def create_eval_steps() -> list:
 
         steps.append(logprobs_step)
 
-    for model, metadata in list(zip(*MARIN_SCALING_SUITES["common_pile"], strict=False)):
-        name = f"marin-comma-{metadata[0]}C-{metadata[-3] * metadata[-2] * 4096}T-{metadata[1]}W-{metadata[2]}D"
+    for model, candidate in list(zip(*MARIN_SCALING_SUITES["common_pile"], strict=False)):
+        total_tokens = int(candidate.tokens)
+        params = candidate.model_config.total_trainable_params(VOCAB_SIZE)
+        name = f"marin-comma-{candidate.flops_budget:.0e}C-{total_tokens}T-N{params:.0e}"
 
         step = evaluate_levanter_lm_evaluation_harness(
             model_name=name,
@@ -73,9 +82,10 @@ def create_eval_steps() -> list:
         )
         steps.append(step)
 
+        model_config = candidate.model_config
         logprobs_step = default_lm_log_probs(
             output_path_of(model).cd("checkpoints"),
-            metadata[-1],
+            model_config,
             dist_eval,
             resource_config=ResourceConfig.with_tpu("v5p-8"),
             checkpoint_is_hf=False,
@@ -84,8 +94,10 @@ def create_eval_steps() -> list:
 
         steps.append(logprobs_step)
 
-    for model, metadata in list(zip(*MARIN_SCALING_SUITES["dclm-default"], strict=False)):
-        name = f"marin-dclm-{metadata[0]}C-{metadata[-3] * metadata[-2] * 4096}T-{metadata[1]}W-{metadata[2]}D"
+    for model, candidate in list(zip(*MARIN_SCALING_SUITES["dclm-default"], strict=False)):
+        total_tokens = int(candidate.tokens)
+        params = candidate.model_config.total_trainable_params(VOCAB_SIZE)
+        name = f"marin-dclm-{candidate.flops_budget:.0e}C-{total_tokens}T-N{params:.0e}"
 
         step = evaluate_levanter_lm_evaluation_harness(
             model_name=name,
@@ -95,16 +107,17 @@ def create_eval_steps() -> list:
         )
         steps.append(step)
 
-    logprobs_step = default_lm_log_probs(
-        output_path_of(model).cd("checkpoints"),
-        metadata[-1],
-        dist_eval,
-        resource_config=ResourceConfig.with_tpu("v5p-8"),
-        checkpoint_is_hf=False,
-        name=versioned(f"{name}-DistRobust-ICE-logprobs"),
-    )
+        model_config = candidate.model_config
+        logprobs_step = default_lm_log_probs(
+            output_path_of(model).cd("checkpoints"),
+            model_config,
+            dist_eval,
+            resource_config=ResourceConfig.with_tpu("v5p-8"),
+            checkpoint_is_hf=False,
+            name=versioned(f"{name}-DistRobust-ICE-logprobs"),
+        )
 
-    steps.append(logprobs_step)
+        steps.append(logprobs_step)
 
     baselines = [
         ("allenai/OLMo-2-1124-7B", "stage2-ingredient3-step8000-tokens34B"),