Scaling Plots & Analysis as an Executor Step

[Helw150]

Closes #2166

[rjpower]

Thanks, I like the idea!

I had a lot of feedback about the structure, which I acknowledge is sort of inherited and therefore you should feel free to push back about how much you want balance telling me I'm wrong, fixing now, or filing issues to handle later.

In general I'd encourage pruning the tests a bit to focus on integration and end-to-end tests which validate a specific output we want to be reproducible (vs constants or "variable is set" tests), and to explicitly try to separate out the handling of wandb vs executor steps vs compute some math.

[eric-czech]

Thanks @Helw150! This is extremely helpful for understanding how to connect all of these things together.

My only somewhat blocking concern is on the assumption that FLOPs from W&B will be on the same scale as FLOPs budgets used in configs. Assuming equal scaling exponents might be another potential issue -- let me know if I'm misunderstanding something there.

[eric-czech]

This seems dangerous to me in a scenario where the config with compute-optimal parameter counts, rather than token counts, did not actually end up being included in the hidden_size loop. I suspect that's more of a problem for experiments like mine where the scaling exponents don't come out equal (and at a limit, models become huge compared to token counts w/ more compute), but either way I think it would be ideal if this code was at least calculating compute-optimal param counts to be used to check that the optimal config found based on optimal token counts is in the ballpark. Or maybe it could just fail if the optimal config ends up being at either boundary as a quicker way to some guardrails?

I suppose $\alpha = \beta$ (equal exponents) is actually a pretty fundamental assumption of all this code, yea? Otherwise, I'm having a hard time seeing why this Chinchilla calculation wouldn't apply where the optimal token count as a function of compute actually scales with the parameter exponent and vise-versa:

cf.

Put somewhat differently, and to check my understanding, that would mean that this function would need to define optimal_tokens using the param scaling exponent because the $C=6ND$ constraint only makes a simultaneous solution necessary when $\alpha \ne \beta$.

[Helw150]

I'm doing the following:

• Find D_opt at each compute budget by fitting parabolas to actual loss values
• Fit D_opt ~ C^b from those empirical minima (from Approach 2 in Chinchilla)
• Given D_opt and C, we can compute N_opt directly as C/(6D), so we don't need to fit alpha.

This should work regardless of beta == alpha.

The screenshot you have there is corresponding to approach 3 I believe (the parametric loss form) rather than the parabola based isoflop form I am using here.

[eric-czech]

Ah, so the alpha being used to define optimal_tokens is actually the same $b$ as in Approach 3, which should be equal to the half the param loss exponent if $\alpha=\beta$ and $\frac{\alpha}{\alpha + \beta}$ otherwise.

On a related note, what do you typically call $\alpha$ and $\beta$ if $a$ and $b$ are the scaling exponents? Loss exponents?

[rjpower]

So this looks much better, thanks!

A few high-level recommendations:

I'd move the "recipe" protocol to isoflop_sweep itself and replace the use of the LlamaConfig with opaque types or protocols as needed. e.g. isoflops might take:

class ModelConfiguration(Protocol):
def num_params()
... whatever ...

class ScalingCallbacks(Protocol):
def build_model_config() ... -> ModelConfiguration
...
(or put the num params callback here as well if you want)

• Consider grugifying even more, and favor copy-pasting the code from isoflop_sweep into the scaling_ladder experiment. Feel free to file an issue if that's a bridge too far, but that's sort of what I'd recommend us to lean towards: experiments are copy paste, library code is agnostic to the model. (If copy-paste is too painful to contemplate, then we consider "how do i refactor this to be more model agnostic so I feel okay importing" etc)

• There's a bit of inherited magic around the schema of the evaluation code. We should try to put the loading logic for the eval results wherever the code that generated it was, or maybe there's a way to share the schema that was used to generate it. (Though it's good you're trying to separate out the "load this specific eval format" from "analysis an eval table").

Overall it's easier to follow for me and less enmeshed with the specifics of Marin/Levanter configuration which is nice!

[rjpower]

Thanks, this is much better.

I'd recommend merging some of the data types and recipes but otherwise I think it's looking good!

[rjpower]

it might be a bit simpler to flatten everything to be part of the "Recipe" protocol, e.g.:

def flops_per_token(candidate: CandidateConfig) -> float
def total_trainable_params(candidate: CandidateConfig) -> float

then the user only needs to define the one protocol

[Helw150]

This one I'll push back on a bit. This protocol defines things that any model config class should have independent of scaling laws (right now in practice our base class is LlamaConfig, which already had these defined).

I don't think things that are features of the model itself in abstract, existed previously to scaling on our base class in practice, and are independently used from scaling should live on the recipe. it'll just force the user to implement needless boilerplate wrapper methods.

[rjpower]

Sure, it's definitely not a big deal. By your argument estimate_memory_bytes should be here, since it seems naturally a function of the model configuration, not the scaling recipe?

[Helw150]

No, because the memory usage in bytes is dependent on both the model configuration and the training specifics (e.g. batch size, which optimizer, etc.). The same model will use different amounts of memory depending on those things.

[eric-czech]

This protocol defines things that any model config class should have independent of scaling laws

+1 to separating model properties from scaling law training configuration like this, and that it makes sense to be able to contextualize "model properties" to specific scaling experiments (i.e. these shouldn't be functions on LlamaConfig). This feels like the right place to me to be able to ablate more controversial aspects of "effective" param counting like embedding inclusion, gated units, weight ties or maybe sparse vs active and quantization scaling params down the road. I can imagine a lot of boilerplate around creating ScalingRecipes otherwise.