Fix po2 for float quant

[Giuseppe5]

Breaking change. Loading old quantized model requires setting a flag for numerical compatibility.
Depends on #1039

The new flow assume that restrict_init_impl (e.g., log2) is done before converting to parameters. After this PR, the op is done after.
This means that loading the model from an old checkpoint would cause log2 to be applied twice.
To avoid this, the a new config flag is created (maybe to extend to a ENV variable as well).
When the flag is set, the reverse op to restrict_init is used (e.g., po2) during checkpoint loading time

[Giuseppe5]

@nickfraser
I tried investigating the LLM test failure but:

• locally, running the test on dev I get a different perplexity
• Running the branch, I get the same perplexity I get locally on dev (but different than the expected one)

[nickfraser]

Let's chat about some bits. Otherwise, looks good.

[nickfraser]

Running the branch, I get the same perplexity I get locally on dev (but different than the expected one)

I'll confirm this on my end too.

The perplexity measurement is a loose guard to force us to do a sanity check like this - assuming I see the same behaviour as you, I'm happy for you to update the scores to make the test pass.

[nickfraser]

Local tests pass. Happy for you to merge when the minor fixes above are corrected.