| """Adapter to expose MegaBlocks package, if available.""" | |
| try: | |
| import megablocks | |
| except ImportError: | |
| megablocks = None | |
| def megablocks_is_available(): | |
| return megablocks is not None | |
| def assert_megablocks_is_available(): | |
| assert ( | |
| megablocks_is_available() | |
| ), "MegaBlocks not available. Please run `pip install megablocks`." | |
| moe = megablocks.layers.moe if megablocks_is_available() else None | |
| dmoe = megablocks.layers.dmoe if megablocks_is_available() else None | |
| arguments = megablocks.layers.arguments if megablocks_is_available() else None | |
| def as_megablocks_args(neox_args): | |
| import copy | |
| tmp = copy.copy(neox_args) | |
| args = arguments.from_megatron(tmp) | |
| args.moe_lbl_in_fp32 = True | |
| args.fp16 = neox_args.precision == "fp16" | |
| args.moe_loss_weight = neox_args.moe_loss_coeff | |
| return args | |