add balance to init

fmoe/megatron/__init__.py

@@ -10,3 +10,10 @@ from .checkpoint import save_checkpoint
 from .checkpoint import load_checkpoint
 
 from .distributed import DistributedDataParallel
+
+from .balance import reset_gate_hook
+from .balance import get_balance_profile
+from .balance import generate_megatron_gate_hook
+from .balance import add_balance_log
+from .balance import patch_forward_step
+from .balance import patch_model_provider