Parallel Scaling Law

import math
from typing import List, Dict

# Coefficients per group for log(loss) = c0 + c1*log(num_params) + c2*log(parallel_size)
_COEFS = {
    'pile': {'c0': 2.0391734962783374, 'c1': -0.06451510348874577, 'c2': -0.018877634923324237},
    'stack': {'c0': 1.4884196746596294, 'c1': -0.0663711422624959, 'c2': -0.02210924328367317},
}

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): 'loss'.
    """
    if group not in _COEFS:
        raise ValueError(f"Unknown group: {group}")
    c0 = _COEFS[group]['c0']
    c1 = _COEFS[group]['c1']
    c2 = _COEFS[group]['c2']
    predictions = []
    for point in input_data:
        num_params = point['num_params']
        parallel_size = point['parallel_size']
        log_loss = c0 + c1 * math.log(num_params) + c2 * math.log(parallel_size)
        loss = math.exp(log_loss)
        predictions.append({'loss': loss})
    return predictions

import math

# Scaling law coefficients per group
_COEFFICIENTS = {
    'stack': {'A': 4.430089002263504, 'b': 0.0663711422624959, 'c': 0.02210924328367317},
    'pile':  {'A': 7.684255507790626, 'b': 0.06451510348874577, 'c': 0.018877634923324237},
}

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

    Args:
        input_data: List of dicts with 'num_params' and 'parallel_size'.
        group: Name of the experimental group.

    Returns:
        List of dicts with predicted 'loss'.
    """
    if group not in _COEFFICIENTS:
        raise ValueError(f'Unknown group: {group}')
    coeffs = _COEFFICIENTS[group]
    results = []
    for item in input_data:
        num_params = item['num_params']
        parallel_size = item['parallel_size']
        loss = coeffs['A'] * (num_params ** (-coeffs['b'])) * (parallel_size ** (-coeffs['c']))
        results.append({'loss': loss})
    return results

import math

_PARAMS = {
    'pile': (7.684255507790626, -0.06451510348874577, -0.018877634923324237),
    'stack': (4.430089002263504, -0.0663711422624959, -0.02210924328367317),
}

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: List of dicts with keys 'num_params' and 'parallel_size'.
        group: Experimental group name ('pile' or 'stack').

    Returns:
        List of dicts, each containing the predicted 'loss'.
    """
    if group not in _PARAMS:
        raise ValueError(f"Unknown group: {group}")
    A, alpha, beta = _PARAMS[group]
    results = []
    for x in input_data:
        N = x['num_params']
        P = x['parallel_size']
        loss = A * (N ** alpha) * (P ** beta)
        results.append({'loss': loss})
    return results

"""
Scaling law prediction function for model parallelism experiments.
"""

from typing import List, Dict

# Pre-fitted coefficients for each experimental group
_PARAMS_MAPPING: Dict[str, Dict[str, float]] = {
    'pile': {
        'A': 7.684255507790626,
        'alpha': -0.06451510348874577,
        'beta': -0.018877634923324237,
    },
    'stack': {
        'A': 4.430089002263504,
        'alpha': -0.0663711422624959,
        'beta': -0.02210924328367317,
    },
}

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 is the same for all groups,
               but the constant parameters/coefficients differ per group.

    Returns:
        A list of dictionaries, corresponding to the input_data list, with each
        dictionary containing the predicted output variable(s) (key 'loss').
    """
    if group not in _PARAMS_MAPPING:
        raise ValueError(f"Unknown group: {group}")
    params = _PARAMS_MAPPING[group]
    A = params['A']
    alpha = params['alpha']
    beta = params['beta']
    predictions: List[Dict[str, float]] = []
    for entry in input_data:
        num_params = entry['num_params']
        parallel_size = entry['parallel_size']
        loss_pred = A * (num_params ** alpha) * (parallel_size ** beta)
        predictions.append({'loss': loss_pred})
    return predictions