Add zloss

[Muennighoff]

No description provided.

[dirkgr]

Anything we can do to get this merged? We want to polish the OLMoE release, and it would be nice if we could just take a dependency on your version instead of @Muennighoff's.

[mvpatel2000]

@Muennighoff can you resolve merge conflicts and make sure tests are passing?

CC: @josejg for review

[josejg]

1. I believe the current implementation will consume extra memory even when z-loss weight is zero. Even though the tensors are small, I always worry about keeping stuff around messing with the mem allocator. When we tested this, we had the zloss tensors separate from the LBL loss using a separate global variable and helper methods, which wouldn't save the logits if the zloss is disabled.

2. I'm also not fan of changing the Router and Experts API (returns and arguments) since that is not backwards compatible. In our implementation, the z-loss was handled entirely in router.py, which prevented this.

3. Similar for changing the API of batched_load_balancing_loss which now returns a tuple instead of a value, not backwards compatible. With our implementation the separate function prevents this.

4. Lastly, we did not observe any meaningful improvements when enabling router zloss, so I would keep it at a default of 0 instead of 1e-3 to prevent the computation in the default case and for backwards compatibility. I see Figure 11 in the OLMoE paper reports stability improvements but I guess the training settings are different.

Here's our implementation for reference - main...josejg/zloss . The main difference is that the modeling code is responsible for doing clear_router_zloss and batched_router_zloss in the corresponding places.

@mvpatel2000 Wdyt?

@Muennighoff could you refactor the PR to address these comments? We can also merge our implementation giving you credit in the PR if that's easier.

[Muennighoff]

oh just merge yours 👍

[josejg]

Sg, tracking it here - #151

[mihir-db]

Closing in favor of #151
@Muennighoff @dirkgr please feel free to comment on other PR if you need any related changes or you can open a github issue :)