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).
    """

    # Fitted parameters for each group
    # Formula: loss = a * num_params^(-alpha) * parallel_size^(-beta) + c
    params = {
        'pile': {
            'a': 5.6297350867e+01,
            'alpha': 0.2142919486,
            'beta': 0.0605768805,
            'c': 1.3518847424
        },
        'stack': {
            'a': 3.7272084315e+01,
            'alpha': 0.2229786185,
            'beta': 0.0719782566,
            'c': 0.7511985021
        }
    }

    # Get parameters for the specified group
    if group not in params:
        raise ValueError(f"Unknown group: {group}. Available groups: {list(params.keys())}")

    p = params[group]
    a = p['a']
    alpha = p['alpha']
    beta = p['beta']
    c = p['c']

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

        # Apply the scaling law
        loss = a * (num_params ** (-alpha)) * (parallel_size ** (-beta)) + 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).
    """
    # Fitted parameters for each group
    # Formula: loss = a * N^(-b) * P^(-d) + c
    # where N = num_params, P = parallel_size

    parameters = {
        'stack': {
            'a': 37.272,
            'b': 0.222979,
            'd': 0.071978,
            'c': 0.751199
        },
        'pile': {
            'a': 56.297,
            'b': 0.214292,
            'd': 0.060577,
            'c': 1.351885
        }
    }

    # Get parameters for the specified group
    if group not in parameters:
        raise ValueError(f"Unknown group: {group}. Available groups: {list(parameters.keys())}")

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

    # Make predictions
    results = []
    for data_point in input_data:
        N = data_point['num_params']
        P = data_point['parallel_size']

        # Apply scaling law: loss = a * N^(-b) * P^(-d) + c
        loss = a * (N ** (-b)) * (P ** (-d)) + 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).
    """

    # Fitted parameters for each group
    # Scaling law form: loss = A * num_params^B * parallel_size^C
    params = {
        'stack': {
            'A': 4.430089,
            'B': -0.066371,
            'C': -0.022109
        },
        'pile': {
            'A': 7.684256,
            'B': -0.064515,
            'C': -0.018878
        }
    }

    # Get parameters for the specified group
    group_params = params[group]
    A = group_params['A']
    B = group_params['B']
    C = group_params['C']

    # Apply the scaling law to each input data point
    predictions = []
    for data_point in input_data:
        num_params = data_point['num_params']
        parallel_size = data_point['parallel_size']

        # Calculate predicted loss using the power law
        # loss = A * num_params^B * parallel_size^C
        predicted_loss = A * (num_params ** B) * (parallel_size ** C)

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

    return predictions

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).
    """
    # Fitted parameters for each group
    # Model: loss = (a - b * parallel_size) * num_params^c
    params = {
        'pile': {
            'a': 7.77714938,
            'b': 0.06308945,
            'c': -0.06479706
        },
        'stack': {
            'a': 4.49228382,
            'b': 0.04233489,
            'c': -0.06670144
        }
    }

    # Get parameters for the specified group
    if group not in params:
        raise ValueError(f"Unknown group: {group}. Available groups: {list(params.keys())}")

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

    # Apply the scaling law to each input data point
    results = []
    for data_point in input_data:
        num_params = data_point['num_params']
        parallel_size = data_point['parallel_size']

        # Compute predicted loss using the scaling law
        # loss = (a - b * P) * N^c
        predicted_loss = (a - b * parallel_size) * (num_params ** c)

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

    return results