fix crash in when running FSDP2+TP (#43226)
* fix crash in when running FSDP2+TP
Signed-off-by: Wang, Yi <yi.a.wang@intel.com>
* update
Signed-off-by: Wang, Yi <yi.a.wang@intel.com>
* fix tp only issue
Signed-off-by: Wang, Yi <yi.a.wang@intel.com>
---------
Signed-off-by: Wang, Yi <yi.a.wang@intel.com>
Co-authored-by: Marc Sun <57196510+SunMarc@users.noreply.github.com>
