Parallel Scaling Law

def law(input_data: list[dict[str, float]], group: str) -> list[dict[str, float]]:
    """
    Predicts output variables based on input variables according to a discovered scaling law.

    Args:
        input_data: A list of dictionaries, where each dictionary is a single data
                    point containing input variable names as keys and their
                    corresponding values.
        group: The name of the experimental group for which to make predictions.
                The functional form of the law must be the same for all groups,
                but the constant parameters/coefficients can differ per group.

    Returns:
        A list of dictionaries, corresponding to the input_data list, with each
        dictionary containing the predicted output variable(s).
    """
    # Parameters fitted from experimental data for each group
    params = {
        'pile': {
            'a': 7.6842555078,
            'b': -0.0645151035,
            'c': -0.0188776349
        },
        'stack': {
            'a': 4.4300890023,
            'b': -0.0663711423,
            'c': -0.0221092433
        }
    }

    if group not in params:
        raise ValueError(f"Unknown group: {group}")

    p = params[group]
    results = []

    for data_point in input_data:
        num_params = data_point['num_params']
        parallel_size = data_point['parallel_size']

        # Scaling law: loss = a * (num_params)^b * (parallel_size)^c
        loss = p['a'] * (num_params ** p['b']) * (parallel_size ** p['c'])

        results.append({'loss': loss})

    return results

def law(input_data: list[dict[str, float]], group: str) -> list[dict[str, float]]:
    """
    Predicts output variables based on input variables according to a discovered scaling law.

    Args:
        input_data: A list of dictionaries, where each dictionary is a single data
                    point containing input variable names as keys and their
                    corresponding values.
        group: The name of the experimental group for which to make predictions.
                The functional form of the law must be the same for all groups,
                but the constant parameters/coefficients can differ per group.

    Returns:
        A list of dictionaries, corresponding to the input_data list, with each
        dictionary containing the predicted output variable(s).
    """
    # Parameters fitted for each group
    # Scaling law: loss = (a0 / parallel_size^alpha) / num_params^beta
    group_params = {
        'pile': {
            'a0': 7.729470668033118,
            'alpha': 0.01889203098571857,
            'beta': 0.06479338533619376
        },
        'stack': {
            'a0': 4.4607706084256336,
            'alpha': 0.022185686139583268,
            'beta': 0.06669649102868953
        }
    }

    if group not in group_params:
        raise ValueError(f"Unknown group: {group}")

    params = group_params[group]
    results = []

    for data_point in input_data:
        num_params = data_point['num_params']
        parallel_size = data_point['parallel_size']

        # Apply the scaling law formula
        loss = (params['a0'] / (parallel_size ** params['alpha'])) / (num_params ** params['beta'])

        results.append({'loss': loss})

    return results

import numpy as np


def law(input_data: list[dict[str, float]], group: str) -> list[dict[str, float]]:
    """
    Predicts output variables based on input variables according to a discovered scaling law.

    Args:
        input_data: A list of dictionaries, where each dictionary is a single data
                    point containing input variable names as keys and their
                    corresponding values.
        group: The name of the experimental group for which to make predictions.
                The functional form of the law must be the same for all groups,
                but the constant parameters/coefficients can differ per group.

    Returns:
        A list of dictionaries, corresponding to the input_data list, with each
        dictionary containing the predicted output variable(s).
    """

    # Parameters fitted for each group
    # Discovered formula: loss = a * num_params^b + c / parallel_size
    params = {
        'stack': {
            'a': 4.438400388168969,
            'b': -0.06832022955222652,
            'c': 0.04365118811808907
        },
        'pile': {
            'a': 7.689659099632834,
            'b': -0.06612727194941805,
            'c': 0.06681072189163119
        }
    }

    if group not in params:
        raise ValueError(f"Unknown group: {group}. Must be one of {list(params.keys())}")

    a = params[group]['a']
    b = params[group]['b']
    c = params[group]['c']

    results = []
    for data_point in input_data:
        num_params = data_point['num_params']
        parallel_size = data_point['parallel_size']

        # Apply the discovered scaling law
        predicted_loss = a * (num_params ** b) + c / parallel_size

        results.append({'loss': predicted_loss})

    return results

def law(input_data: list[dict[str, float]], group: str) -> list[dict[str, float]]:
    """
    Predicts output variables based on input variables according to a discovered scaling law.

    The scaling law models the final language modeling loss as a function of the number of
    model parameters and the degree of model parallelism. The discovered law follows the form:

        loss = (c - delta * ln(parallel_size)) * num_params^b

    where:
    - c and b are group-specific coefficients
    - delta is a universal coefficient controlling the parallelism benefit
    - num_params is the total number of model parameters
    - parallel_size is the degree of parallelism (1, 2, 4, etc.)

    Args:
        input_data: A list of dictionaries, where each dictionary is a single data
                    point containing input variable names as keys and their
                    corresponding values. Expected keys: 'num_params', 'parallel_size'
        group: The name of the experimental group for which to make predictions.
               Valid values: 'stack', 'pile'

    Returns:
        A list of dictionaries, corresponding to the input_data list, with each
        dictionary containing the predicted output variable 'loss'.
    """
    import math

    # Group-specific coefficients fitted from training data
    params = {
        'stack': {'c': 4.639114, 'b': -0.068405},
        'pile': {'c': 7.632338, 'b': -0.064275}
    }

    # Universal coefficient for parallelism benefit
    delta = 0.123483

    # Validate group parameter
    if group not in params:
        raise ValueError(f"Unknown group: {group}. Valid groups are: {list(params.keys())}")

    # Get group-specific parameters
    c = params[group]['c']
    b = params[group]['b']

    # Make predictions for each input data point
    results = []
    for data_point in input_data:
        num_params = float(data_point['num_params'])
        parallel_size = float(data_point['parallel_size'])

        # Apply the scaling law formula
        # loss = (c - delta * ln(parallel_size)) * num_params^b
        coefficient = c - delta * math.log(parallel_size)
        loss = coefficient * (num_params ** b)

        results.append({'loss': loss})

    return results