selective compilation - norm layers only

[lessw2020]

This PR adds the option to selectively compile just the norm layers only, and is mainly targeted at RMSNorm.
By compiling just the norm layers when using rmsnorm, we get nearly comparable speedups as using the fusedRMSNorm triton kernel.
Credit @wconstab for this idea.

regular rmsnorm:

with the new compile_rmsnorm enabled:

2 - UX - I enabled the compile rmsnorm as it's own option for now so users can quickly try whole model or norm only compile. If compile is true, then the rmsnorm layers will not specifically be compiled (as they will be included in the generic full model compile) and a minor note is issued in logging.

3 - using other norms with this option enabled does not appear to add any speedup (but also no errors) so I did not add a check to only compile if norm is rmsnorm (but can add that).

[drisspg]

Is 2 saying that in order to have "full" compile you need to set both compile=true and compile_rmsnorm = true

[drisspg]

fullgraph?

[tianyu-l]

1. I think we should not couple compile-related code with parallelize_llama code.
   Maybe we should create a file called compile.py and put everything there. This was not viable when we need to do per TransformerBlock compile due to some restriction (vaguely remember we need to first compile then FSDP wrap?), but now there's no such restriction since we are compiling the whole model.

2. Would you please also show the comparison between triton kernel and compiled rmsnorm? That would give us a better idea of how much the improvement is.

3. Let's follow up on if it works with TP, FSDP+TP, PP, 3D, SAC, etc.

[tianyu-l]

just calling it compile_rmsnorm is good enough I think

[tianyu-l]

I think it's better to write compile_rmsnorm = job_config.model.norm_type == "rmsnorm" and job_config.training.compile_rmsnorm and not job_config.training.compile

[tianyu-l]

IMO there's no need to print this message. It's still compiled, although with the whole model.

[tianyu-l]

ditto no need to set it here

[tianyu-l]

If it's called compile_rmsnorm, we should only compile if the norm_type = "rmsnorm". The help message needs to reflect this.

[tianyu-l]

I think we shouldn't enable it everywhere before we understand when it works and when it doens't, e.g. 2D, selective AC, etc.

[lessw2020]

Is 2 saying that in order to have "full" compile you need to set both compile=true and compile_rmsnorm = true

I updated the text to be more specific, but no - if compile = true in the config, then you get full compile including the rmsnorm layers.

[wanchaol]

IMO we should not add another option. We can just possibly reuse the fused_rmsnorm field in the norm_type, instead of adding a separate option here

[wanchaol]

I think we should also check the fqn, iirc torch.compile by default would prepend a fqn for the compiled model, we should try to get rid of it in some way, otherwise we would not be able to load Llama3 pretrained weights

[lessw2020]

you are right - it inserts into the fqn (_orig_mod):

[rank0]:Named parameters of the module:
[rank0]:Parameter name: attention.wq.weight, Shape: torch.Size([256, 256])
[rank0]:Parameter name: attention.wk.weight, Shape: torch.Size([256, 256])
[rank0]:Parameter name: attention.wv.weight, Shape: torch.Size([256, 256])
[rank0]:Parameter name: attention.wo.weight, Shape: torch.Size([256, 256])
[rank0]:Parameter name: feed_forward.w1.weight, Shape: torch.Size([768, 256])
[rank0]:Parameter name: feed_forward.w2.weight, Shape: torch.Size([256, 768])
[rank0]:Parameter name: feed_forward.w3.weight, Shape: torch.Size([768, 256])
[rank0]:Parameter name: attention_norm._orig_mod.weight, Shape: torch.Size([256])
[rank0]:Parameter name: ffn_norm._orig_mod.weight, Shape: torch.Size([256])

[lessw2020]

will make a cleanup function

[wanchaol]

IMO we should probably not make a state_dict hook if that's what u are going to do.

Instead we should compile the norm function directly

[fegin]

distributed_state_dict already has this covered. It will remove the FQN from torch.compile. Let me know if you do see some issues.

[fegin]

An alternative, if you are still seeing this error, is to call module.compile instead of torch.compile

[fegin]

This is actually not correct. If AC is enabled, you will see 2 extra submodules in the AC module, attention_norm, ffn_norm. And these 2 submodules will never be used in the forward but only be used in the state_dict().

[tianyu-l]

close as we removed the feature in #535