Bug Fix - Fix bugs in optional tokens per expert logging (#54)

This PR fixes two bugs in #41 
- `--moe-tokens-logging` doesn't control whether the related tracked names are added.
- `sinkhorn` tokens per expert metrics are never tracked in original if statement.

Author: yzygitzh

megatron/training/training.py

@@ -1527,6 +1527,7 @@ def training_log(loss_dict, total_loss_dict, learning_rate, decoupled_learning_r
         track_names = []
         if args.moe_router_load_balancing_type in ["aux_loss", "seq_aux_loss", "global_batch_loss"]:
             track_names.append("load_balancing_loss")
+        if args.moe_tokens_logging:
             if args.moe_router_load_balancing_type == "sinkhorn":
                 track_names.append("sinkhorn_tokens_per_expert")
             elif args.moe_router_load_balancing_type == "aux_loss":