feat: online load/save huggingface model weights for Megatron-FSDP

[conver334]

What does this PR do ?

Adds support for online loading and saving of Hugging Face model weights for Megatron-FSDP with expert parallelism (EP).

Tensor parallelism (TP) is not supported at this time.

API

see detailed usage in file: example/5.mfsdp_load_and_export_multiple_gpus.py

# Init random Megatron-FSDP model
bridge = AutoBridge.from_pretrained(hf_model_path)
ddp_config = {
    "use_distributed_optimizer": True,
    "check_for_nan_in_grad": True,
    "use_megatron_fsdp": True,
    "data_parallel_sharding_strategy": "optim_grads_params",
}
model = bridge.get_model(wrap_with_ddp=True, use_megatron_fsdp=True, ddp_config=ddp_config,data_parallel_random_init=False, post_model_creation_callbacks=[])
 
# load HF weights
bridge.load_weights(model, hf_model_path, memory_efficient=True)

# export HF weights
for k, v in bridge.export_weights(model): 
    pass

Test

Successfully tested with Qwen2 and DeepSeekV3 architectures.

torchrun --nproc_per_node=8 5.mfsdp_load_and_export_multiple_gpus.py --model_path moonshotai/Moonlight-16B-A3B --ep 4 --trust_remote_code 

@ISEEKYAN