@dataclass
class BuildDataArgs:
model_type_ = None
model_provider_ = None
train_valid_test_datasets_provider_ = None
def set_build_data_args(model_type, model_provider, train_valid_test_datasets_provider):
BuildDataArgs.model_type_ = model_type
BuildDataArgs.model_provider_ = model_provider
BuildDataArgs.train_valid_test_datasets_provider_ = train_valid_test_datasets_provider
以Megatron为例，获取Megatron框架中用以建立模型优化器及数据集的必要参数，客户需根据自身框架决定获取的必要参数。
def save_callback(step: int, save_info: list, train_args=None, ctx=None):
model = train_args[TRAIN_PARAM][MODEL_INDEX]
optimizer = train_args[TRAIN_PARAM][OPTIM_INDEX]
opt_param_scheduler = train_args[TRAIN_PARAM][SCHEDULER_INDEX]
global_args = get_args()
# Update learning rate.
if global_args.train_samples is None:
global_args.consumed_train_samples = step * global_args.global_batch_size
if train_args[TRAIN_PARAM][SCHEDULER_INDEX].num_steps != global_args.consumed_train_samples:
train_args[TRAIN_PARAM][SCHEDULER_INDEX].step(global_args.global_batch_size)
if hasattr(optimizer, 'optim_nums'):
for _, info_dict in enumerate(save_info):
optim_idx = info_dict.get("type", 0)
rank_list = info_dict.get("ranks", None)
save_rank = rank_list[0]
optimizer.set_dump_args(optim_idx, save_rank, step, rank_list)
else:
rank_list = save_info[0].get("ranks", None)
save_rank = rank_list[0]
optimizer.set_dump_args(save_rank, step, rank_list)
save_checkpoint(step, model, optimizer, opt_param_scheduler,
global_args.num_floating_point_operations_so_far)
def rename_callback(step: int, ctx=None):
iteration = step
rank = torch.distributed.get_rank()
args = get_args()
tmp_dir = 'iter_{:07d}_tmp'.format(iteration)
fin_dir = 'iter_{:07d}'.format(iteration)
src_path = os.path.join(args.save, tmp_dir)
dst_path = os.path.join(args.save, fin_dir)
try:
src_check_ret, err_msg, src_abs_path = FileUtils.regular_file_path(src_path, args.save, False)
dst_check_ret, err_msg, dst_abs_path = FileUtils.regular_file_path(dst_path, args.save, False)
if (not src_check_ret) or (not dst_check_ret):
raise FileNotFoundError(f"rank:{rank} rename path error. {err_msg}")
os.rename(src_abs_path, dst_abs_path)
except FileNotFoundError as e:
print(f"{e}")
return False
# And update the latest iteration
tracker_filename = get_checkpoint_tracker_filename(args.save)
is_path_valid, err_msg, tracker_filename = FileUtils.regular_file_path(tracker_filename, "/", False)
if not is_path_valid:
print_rank_0(err_msg)
raise Exception(" tracker_filename is not valid")
with open(tracker_filename, 'w') as f:
f.write(str(iteration))
return True
def get_checkpoint_name(checkpoints_path, iteration, release=False,
pipeline_parallel=None,
tensor_rank=None, pipeline_rank=None,
expert_parallel=None, expert_rank=None,
return_base_dir=False, tmp=False):
if release:
directory = 'release'
else:
directory = 'iter_{:07d}'.format(iteration)
if tmp:
directory = directory + "_tmp"
if return_base_dir:
common_path = os.path.join(checkpoints_path, directory)
return common_path
# Use both the tensor and pipeline MP rank.
if pipeline_parallel is None:
pipeline_parallel = (mpu.get_pipeline_model_parallel_world_size() > 1)
if tensor_rank is None:
tensor_rank = mpu.get_tensor_model_parallel_rank()
if pipeline_rank is None:
pipeline_rank = mpu.get_pipeline_model_parallel_rank()
if expert_parallel is None:
expert_parallel = (mpu.get_expert_model_parallel_world_size() > 1)
if expert_rank is None:
expert_rank = mpu.get_expert_model_parallel_rank()
# Use both the tensor and pipeline MP rank. If using the distributed
# optimizer, then the optimizer's path must additionally include the
# data parallel rank.
if not pipeline_parallel:
common_path = os.path.join(checkpoints_path, directory,
f'mp_rank_{tensor_rank:02d}')
else:
common_path = os.path.join(checkpoints_path, directory,
f'mp_rank_{tensor_rank:02d}_{pipeline_rank:03d}')
if expert_parallel:
common_path = common_path + f'_{expert_rank:03d}'
return os.path.join(common_path, "model_optim_rng.pt")
def save_checkpoint(iteration, model, optimizer, opt_param_scheduler,
num_floating_point_operations_so_far):
"""Save a model checkpoint."""
args = get_args()
# Only rank zero of the data parallel writes to the disk.
model = unwrap_model(model)
ckpt_format = args.dist_ckpt_format if args.use_dist_ckpt else 'torch'
print('rank {} is saving checkpoint at iteration {:7d} to {} in {} format'
.format(args.rank, iteration, args.save, ckpt_format))
# Collect rng state across data parallel ranks.
rng_state = get_rng_state(args.use_dist_ckpt)
tmp = optimizer.get_error_dump()
# Checkpoint name.
checkpoint_name = get_checkpoint_name(args.save, iteration, return_base_dir=args.use_dist_ckpt, tmp=tmp)
check_ret, err_msg, checkpoint_name = FileUtils.regular_file_path(checkpoint_name, '/', False)
if not check_ret:
raise Exception(f"rank {args.rank} get checkpoint name error, {err_msg}")
# Save distributed optimizer's custom parameter state.
if args.use_distributed_optimizer and not args.no_save_optim and optimizer is not None and not args.use_dist_ckpt:
optim_checkpoint_name = \
get_distributed_optimizer_checkpoint_name(checkpoint_name)
ensure_directory_exists(optim_checkpoint_name)
optimizer.save_parameter_state(optim_checkpoint_name)
if not args.use_distributed_optimizer and (args.fp16 or args.bf16):
if not hasattr(optimizer.config, 'reuse_fp32_param'):
optimizer._copy_main_params_to_model_params()
cur_rank = torch.distributed.get_rank()
save_flag = optimizer.need_write_file()
# Collect args, model, RNG.
if not torch.distributed.is_initialized() \
or save_flag \
or args.use_dist_ckpt:
optim_sd_kwargs = {}
if args.use_dist_ckpt and args.use_distributed_optimizer:
optim_sd_kwargs['sharding_type'] = ('fully_sharded_bucket_space'
if args.ckpt_fully_parallel_save
else 'dp_zero_gather_scatter')
print_rank_0(f'Storing distributed optimizer sharded state of type {optim_sd_kwargs["sharding_type"]}')
state_dict = generate_state_dict(args, model, optimizer, opt_param_scheduler, rng_state,
args.use_dist_ckpt, iteration, optim_sd_kwargs=optim_sd_kwargs)
state_dict['num_floating_point_operations_so_far'] = num_floating_point_operations_so_far
# Save.
ensure_directory_exists(checkpoint_name)
torch.save(state_dict, checkpoint_name)
MindIO TTP功能沿用了Megatron框架中保存CheckPoint的主要逻辑，但是由于需要参与CheckPoint保存的worker不再是有序的，因此需要对CheckPoint文件的命名以及是否进行CheckPoint保存的判断逻辑进行修改，在临终保存之后文件名可能与Megatron框架中保存CheckPoint的文件名不同，由rename_callback实现文件改名以和框架中的命名规则一致。
def stop_callback(train_args=None, ctx=None):
# stop and clean device
device = get_current_device()
torch.npu.set_device(device)
torch_npu.npu.stop_device(device)
def clean_callback(is_uce_error=False, train_args=None, ctx=None):
"""
this function do:
1) get UCE check result from torch_npu
2) do some clear before rebuild (avoid OOM) when the check result is UCE_HIGH_LEVEL
3) HCCL resume and restart device
"""
device = get_device()
rank = torch.distributed.get_rank()
torch.npu.set_device(device)
ret = RET_OK
if is_uce_error:
check_memory_result = torch_npu.npu.check_uce_in_memory(device)
if check_memory_result == UCE_LOW_LEVEL: # no need rebuild
ret = RET_NO_REBUILD
elif check_memory_result == UCE_HIGH_LEVEL: # need rebuild
clear_memory(train_args[TRAIN_PARAM][MODEL_INDEX],
train_args[TRAIN_PARAM][OPTIM_INDEX],
train_args[TRAIN_PARAM][CONFIG_INDEX])
train_args[TRAIN_PARAM][MODEL_INDEX] = None
train_args[TRAIN_PARAM][OPTIM_INDEX] = None
train_args[TRAIN_PARAM][SCHEDULER_INDEX] = None
train_args[TRAIN_PARAM][CONFIG_INDEX] = None
ret = RET_OK
else: # exit
ret = RET_ERROR
if hasattr(torch.distributed, 'reinit_process_group'):
torch.distributed.reinit_process_group(group=None, rebuild_link=False)
torch.npu.restart_device(device)
return ret
def repair_callback(step: int, need_rebuild: bool, error_ranks: list, repair_info: dict,
train_args=None, ctx=None):
from mindio_ttp.framework_ttp import RepairType, OptimizerType
torch.npu.set_device(get_current_device())
optim_idxs = repair_info.get("type", OptimizerType.ATTENTION.value)
repair_type = repair_info.get("repair_type", None)
src_ranks = repair_info.get("src", [])
dest_ranks = repair_info.get("dst", [])
rank_list = repair_info.get("rank_list", [])
build_repair_group(rank_list)
if repair_type == RepairType.RT_SEND.value: # send
return send_rank_repair(src_ranks, dest_ranks, optim_idxs, step, rank_list, train_args[TRAIN_PARAM])
elif repair_type in [RepairType.RT_UCE_HIGHLEVEL.value,
RepairType.RT_UCE_LOWLEVEL.value,
RepairType.RT_RECV_REPAIR.value]:
return recv_rank_repair(src_ranks, dest_ranks, optim_idxs, step,
need_rebuild, rank_list, train_args[TRAIN_PARAM])
else:
return RET_ERROR
def rollback_callback(step: int, train_args=None, ctx=None):
args = get_args()
torch.npu.set_device(get_current_device())
global temp_memory_ckpt
rank = torch.distributed.get_rank()
# Update consumed train samples and learning rate.
if args.train_samples is None:
args.consumed_train_samples = step * args.global_batch_size
if train_args[TRAIN_PARAM][SCHEDULER_INDEX].num_steps != args.consumed_train_samples:
train_args[TRAIN_PARAM][SCHEDULER_INDEX].step(args.global_batch_size)
feature_rollback()
gather_model_params_from_optimizer(train_args[TRAIN_PARAM][OPTIM_INDEX], step)
rebuild_dataset(train_args)
set_memory_ckpt(None)
temp_memory_ckpt = None
destroy_repair_group()
return RET_OK
def feature_rollback():
# Other features re-initialize or roll back data after the fault is rectified.
if hasattr(args, "num_experts") and args.num_experts:
mpu._MOE_AUX_LOSSES_LOGGING_TRACKER = {} # 重新初始化全局Tensor
def send_rank_repair(src_ranks: list, dest_ranks: list, optim_idxs: list,
step: int, rank_list: list, train_args):
rank_list.sort()
rank_ = torch.distributed.get_rank()
repair_group = get_repair_group()
for idx, src_rank in enumerate(src_ranks):
dest_rank, optim_idx = dest_ranks[idx], optim_idxs[idx]
if src_rank != rank_:
return RET_ERROR
if src_rank == dest_rank:
return RET_ERROR
save_and_send_ckpt(src_rank, dest_rank, step, optim_idx, train_args)
for idx, _ in enumerate(src_ranks):
dest_rank, optim_idx = dest_ranks[idx], optim_idxs[idx]
train_args[OPTIM_INDEX].send_optim_param_state(dest_rank, repair_group, step, optim_idx)
# repair other global tensor
for idx, src_rank in enumerate(src_ranks):
torch.distributed.send(global_tensor, dst=dest_rank, group=repair_group)
return RET_OK
def save_and_send_ckpt(src_rank, dest_rank, step, optim_idx, train_args):
"""
Save memory checkpoint and send to dest rank.
"""
repair_group = get_repair_group()
state_dict = save_memory_ckpt(train_args[OPTIM_INDEX], train_args[SCHEDULER_INDEX],
step, optim_idx)
buffer = io.BytesIO()
torch.save(state_dict, buffer)
state_dict_bytes = buffer.getvalue()
state_dict_tensor = torch.ByteTensor(torch.ByteStorage.from_buffer(state_dict_bytes)).to('npu')
# Send tensor size first
size_tensor = torch.tensor([state_dict_tensor.numel()], dtype=torch.long).to('npu')
torch.distributed.send(size_tensor, dst=dest_rank, group=repair_group)
# Send the serialized state_dict tensor
torch.distributed.send(state_dict_tensor, dst=dest_rank, group=repair_group)
return RET_OK
def recv_rank_repair(src_ranks: list, dest_ranks: list, optim_idxs: list,
step: int, need_rebuild: bool, rank_list: list, train_args):
# low level and only rollback case
if src_ranks == dest_ranks:
return RET_OK
rank_ = torch.distributed.get_rank()
rank_list.sort()
if need_rebuild:
build_local_embedding_group()
model, optimizer, lr_scheduler, config = rebuild_model_and_optimizer(
BuildDataArgs.model_provider_, BuildDataArgs.model_type_)
train_args[MODEL_INDEX] = model
train_args[OPTIM_INDEX] = optimizer
train_args[SCHEDULER_INDEX] = lr_scheduler
train_args[CONFIG_INDEX] = config
for idx, src_rank in enumerate(src_ranks):
dest_rank, optim_idx = dest_ranks[idx], optim_idxs[idx]
if dest_rank != rank_:
return RET_ERROR
if src_rank != dest_rank:
recv_ckpt_from_peer(src_rank, dest_rank, step, rank_list)
# combine state_dict and once load,fix precision problem
state_dict = get_memory_ckpt()
load_memory_ckpt(train_args[MODEL_INDEX], train_args[OPTIM_INDEX], train_args[SCHEDULER_INDEX],
state_dict, None)
repair_group = get_repair_group()
for idx, src_rank in enumerate(src_ranks):
dest_rank, optim_idx = dest_ranks[idx], optim_idxs[idx]
train_args[OPTIM_INDEX].recv_and_load_optim_param_state(src_rank, repair_group, step, optim_idx)
# repair other global tensor
for idx, src_rank in enumerate(src_ranks):
recv_tensor = torch.empty(size, dtype=type, device="npu")
torch.distributed.recv(recv_tensor, src=src_rank, group=repair_group)
global_tesor.data.copy_(recv_tensor)
return RET_OK
def recv_ckpt_from_peer(src_rank, dest_rank, step, rank_list: list):
"""
receive memory checkpoint and repair train() param.
"""
repair_group = get_repair_group()
# Receive tensor size first
size_tensor = torch.tensor([0], dtype=torch.long, device='npu')
torch.distributed.recv(size_tensor, src=src_rank, group=repair_group)
size = size_tensor.item()
# Receive the serialized state_dict tensor
state_dict_tensor = torch.empty(size, dtype=torch.uint8, device='npu')
torch.distributed.recv(state_dict_tensor, src=src_rank, group=repair_group)
# Deserialize the state_dict
state_dict_bytes = state_dict_tensor.cpu().numpy().tobytes()
buffer = io.BytesIO(state_dict_bytes)
device_count = torch.npu.device_count()
if device_count == 0:
raise ValueError(f"torch.npu.device_count return 0!")
map_location = {
'npu:' + str(src_rank % device_count): 'npu:' + str(dest_rank % device_count)
}
loaded_state_dict = torch.load(buffer, map_location=map_location)
set_memory_ckpt(loaded_state_dict)
def rebuild_model_and_optimizer(model_provider, model_type):
args = get_args()
load = args.load
args.load = None
ori_embedding_group = mpu._EMBEDDING_GROUP
mpu._EMBEDDING_GROUP = get_local_embedding_group()
model, optimizer, lr_scheduler = setup_model_and_optimizer(model_provider, model_type)
mpu._EMBEDDING_GROUP = ori_embedding_group
args.load = load
config = get_model_config(model[0])
return model, optimizer, lr_scheduler, config
def rebuild_dataset(args):
# repair data iterator
train_data_iterator, valid_data_iterator, test_data_iterator \
= build_data_iterator(args[TRAIN_PARAM][MODEL_INDEX])
args[TRAIN_PARAM][TRAIN_DATA_INDEX] = train_data_iterator
args[TRAIN_PARAM][VALID_DATA_INDEX] = valid_data_iterator
args[TEST_DATA_ITER][0] = test_data_iterator
return RET_OK
def build_data_iterator(model):
args = get_args()
if args.virtual_pipeline_model_parallel_size is not None:
train_data_iterator = []
valid_data_iterator = []
test_data_iterator = []
for i in range(len(model)):
mpu.set_virtual_pipeline_model_parallel_rank(i)
iterators = build_train_valid_test_data_iterators(
BuildDataArgs.train_valid_test_datasets_provider_)
train_data_iterator.append(iterators[0])
valid_data_iterator.append(iterators[1])
test_data_iterator.append(iterators[2])
else:
train_data_iterator, valid_data_iterator, test_data_iterator \
= build_train_valid_test_data_iterators(
BuildDataArgs.train_valid_test_datasets_provider_)
return train_data_iterator, valid_data_iterator, test_data_iterator
def update_arf_reboot_flag(new_state): # 对于MindIO ARF中增量启动的节点，选择修复过程中接收的优化器数据而不是加载故障前已保存的CheckPoint数据
global ARF_REBOOT_FLAG, LOAD, PRETRAINED_CHECKPOINT
args = get_args()
if new_state:
LOAD = args.load
PRETRAINED_CHECKPOINT = args.pretrained_checkpoint
args.load = None
args.pretrained_checkpoint = None
elif ARF_REBOOT_FLAG:
args.load = LOAD
args.pretrained_checkpoint = PRETRAINED_CHECKPOINT
ARF_REBOOT_FLAG = new_state
def arf_rebuild_process_group_callback(fault_ranks: list, train_args, ctx):
models, optimizer = train_args[TRAIN_PARAM][MODEL_INDEX], train_args[TRAIN_PARAM][OPTIM_INDEX]
args = get_args()
update_arf_reboot_flag(False)
# 1.1 destroy_all_process_group
torch.distributed.destroy_process_group()
# 1.2 init_all_process_group
torch.distributed.init_process_group(
backend=args.distributed_backend,
world_size=args.world_size,
rank=args.rank,
timeout=timedelta(minutes=args.distributed_timeout_minutes),
)
#2.1 destroy mpu group
delete_mpu_group(mpu)
# 2.2 destroy global memory buffer
mpu.destroy_global_memory_buffer()
# 3.1 init_mpu
mpu.initialize_model_parallel(
context_parallel_size=args.context_parallel_size,
tensor_model_parallel_size=args.tensor_model_parallel_size,
expert_model_parallel_size=args.expert_model_parallel_size,
distributed_timeout_minutes=args.distributed_timeout_minutes,
pipeline_model_parallel_size=args.pipeline_model_parallel_size,
nccl_communicator_config_path=args.nccl_communicator_config_path,
pipeline_model_parallel_split_rank=args.pipeline_model_parallel_split_rank,
virtual_pipeline_model_parallel_size=args.virtual_pipeline_model_parallel_size,
)
update_model_and_optim_related_group(models, optimizer)
def delete_mpu_group(mpu_module):
group_space = [mpu_group for mpu_group in dir(mpu_module) if 'GROUP' in mpu_group]
for group_name in group_space:
setattr(mpu_module, group_name, None)
def update_model_and_optim_related_group(models, optimizer):
for model in models:
# dense
for buffer in model.buffers:
# fix ParamAndGradBuffer attributes
buffer.data_parallel_group = mpu._DATA_PARALLEL_GROUP_WITH_CP
buffer.data_parallel_world_size = torch.distributed.get_world_size(group=mpu._DATA_PARALLEL_GROUP_WITH_CP)
for bucket in buffer.buckets:
# fix Bucket attributes
bucket.data_parallel_group = mpu._DATA_PARALLEL_GROUP_WITH_CP
bucket.data_parallel_world_size = torch.distributed.get_world_size(
group=mpu._DATA_PARALLEL_GROUP_WITH_CP)
bucket.data_parallel_rank = torch.distributed.get_rank(mpu._DATA_PARALLEL_GROUP_WITH_CP)
# moe
for expert_buffer in model.expert_parallel_buffers:
# fix ParamAndGradBuffer attributes
expert_buffer.data_parallel_group = mpu._DATA_MODULO_EXPERT_PARALLEL_GROUP
expert_buffer.data_parallel_world_size = \
torch.distributed.get_world_size(group=mpu._DATA_MODULO_EXPERT_PARALLEL_GROUP)
for bucket in expert_buffer.buckets:
# fix Bucket attributes
bucket.data_parallel_group = mpu._DATA_MODULO_EXPERT_PARALLEL_GROUP
bucket.data_parallel_world_size = torch.distributed.get_world_size(
group=mpu._DATA_MODULO_EXPERT_PARALLEL_GROUP)
bucket.data_parallel_rank = torch.distributed.get_rank(mpu._DATA_MODULO_EXPERT_PARALLEL_GROUP)
if hasattr(optimizer, 'optim_nums'): # 根据需求更新优化器中持有的通信组
optimizer.chained_optimizers[0].ori_dp_group = mpu._DATA_PARALLEL_GROUP_WITH_CP
optimizer.chained_optimizers[1].ori_dp_group = mpu._DATA_MODULO_EXPERT_PARALLEL_GROUP
else:
optimizer.ori_dp_group = mpu._DATA_PARALLEL_GROUP_WITH_CP
if not get_args().use_distributed_optimizer: # 非分布式优化器不持有副本通信组
return
# fix optimizer attributes
if hasattr(optimizer, 'optim_nums'): # 分布式优化器会持有副本组用于副本优化器相关的通信
optimizer.chained_optimizers[0].data_parallel_group = ttp_get_dp_cp_replica_group()
optimizer.chained_optimizers[0].data_parallel_group_gloo = ttp_get_dp_cp_replica_group_gloo()
optimizer.chained_optimizers[0].ori_dp_list = torch.distributed.get_process_group_ranks(
mpu._DATA_PARALLEL_GROUP_WITH_CP)
optimizer.chained_optimizers[1].data_parallel_group = ttp_get_dp_ep_replica_group()
optimizer.chained_optimizers[1].data_parallel_group_gloo = ttp_get_dp_ep_replica_group_gloo()
optimizer.chained_optimizers[1].ori_dp_list = torch.distributed.get_process_group_ranks(
mpu._DATA_MODULO_EXPERT_PARALLEL_GROUP)
else:
optimizer.data_parallel_group = ttp_get_dp_cp_replica_group()
optimizer.data_parallel_group_gloo = ttp_get_dp_cp_replica_group_gloo()
optimizer.ori_dp_list = torch.distributed.get_process_group_ranks(mpu._DATA_PARALLEL_GROUP_WITH_CP)
return
#（可选）若训练中涉及切流操作，重置训练中使用的流，确保调度正确
class CustomFunction(torch.autograd.Function):
@staticmethod
def forward(ctx, input):
torch.cuda.set_stream(torch.cuda.default_stream())
return input
@staticmethod
def backward(ctx, grad):
torch.cuda.set_stream(torch.cuda.default_stream())
return grad
def recover_set_stream():
input_tensor = torch.empty(1, dtype=torch.float32, device="npu", requires_grad=True)
grad_tensor = torch.empty(1, dtype=torch.float32, device="npu", requires_grad=True)
output_tensor = CustomFunction.apply(input_tensor)
output_tensor.backward(grad_tensor)
stop&clean callback为MindIO UCE和MindIO ARF必要的回调函数，stop_callback用于通过torch_npu.npu.stop_device接口令其他worker暂停训练，待修复完成后所有worker共同重新进入训练；clean_callback中，如果当前故障为UCE故障，则需要通过torch_npu.npu.check_uce_in_memory接口检测UCE故障类型，若为UCE_HIGH_LEVEL类型，则本次修复需要重建模型优化器以重新申请内存，刷新Tensor状态以完全修复UCE故障；之后通过torch.npu.restart_device接口重新启动device。
repair_callback用于实现互为副本的worker进行数据修复，故障worker重建模型优化器并接收并加载正常副本worker的CheckPoint数据，正常副本worker保存并发送CheckPoint数据。
在repair完成后，所有worker共同进入rollback_callback中，在rollback阶段，将优化器中的参数取出到模型，还需要将其他训练中用到的参数回滚至正确的位置，如学习率、数据集。
MindIO ARF功能依赖MindIO UCE修复中用到的数据传输功能，arf_rebuild_process_group_callback用于MindIO ARF功能在增量重启worker后重建通信组，并替换模型优化器中持有的通信组，确保在故障修复后的训练中使用重建的通信组。
此外，对于其他需要修复的全局Tensor，若其与训练迭代无依赖关系，则可在rollback阶段对其重新初始化，将其释放；若其与训练迭代有依赖关系，且与副本优化器映射关系一致时，则需要在repair阶段对其进行重建并通过点对点通信修复该Tensor的数据。