feat: introduce llama4 support

[ashwinb]

As the title says. Details in README.

[awgu]

Even though 2b2e5b2 tried to clean this docstring up, the current state in main is not correct:

llama-models/models/llama4/moe.py

Lines 104 to 118 in 63172b3

	Tensors used in this module are annotated with the suffixes that indicate the shape of the tensor.
	Several commonly used annotations include:
	- a: bsz*slen
	- E: number of experts
	- e: number of local experts per ep (n_experts/ep)
	- D: hidden dimension
	- d: D/tp
	- F: model dimension
	- G: number of tokens per expert (a * capacity_factor / E)
	- g: number of tokens per expert per TP rank (i.e., G/TP)
	Examples:
	x_aD [a, D]
	routed_in_etG_D [et*G, D]
	x_eGD: [e, G, D]

1. I think that capacity_factor and top_k in MoEArgs should have been combined to be the same thing. The comment for capacity_factor is incorrectly for EC, not TC.
   

   llama-models/models/llama4/args.py

   Lines 29 to 35 in 63172b3

   	class MoEArgs(BaseModel):
   	num_experts: int = -1
   	capacity_factor: float = 1.0 # capacity factor determines how many tokens each expert can choose
   	auto_scale_F: bool = ( # noqa: N815
   	True # if true, rescales hidden_dim such that number of activated params is same as equivalent dense layer
   	)
   	top_k: int = 1

2. The annotations G and g as described are for EC, not TC. This means that annotations like routed_in_EG_D are incorrect. The token dimension is not sharded coming out of the experts forward, so the annotation routed_out_eg_D is incorrect.
   

   llama-models/models/llama4/moe.py

   Line 201 in 63172b3

   routed_out_eg_D = self.experts(routed_in_EG_D.detach())

Since this is the reference implementation, I hope that some more care can be taken in polishing it. Previously, repos like torchtitan used the reference impl as the starting point for their own model definitions, and others in open source may read it to learn from as well. Thanks!