source:Avoid redundant transpose of fused weights when weight_loader is called iteratively
(
self,
param,
loaded_weight,
expert_id,
shard_id: Optional[str] = None,
source: Optional[str] = None,
)
| 256 | ) |
| 257 | |
| 258 | def weight_loader( |
| 259 | self, |
| 260 | param, |
| 261 | loaded_weight, |
| 262 | expert_id, |
| 263 | shard_id: Optional[str] = None, |
| 264 | source: Optional[str] = None, |
| 265 | ): |
| 266 | """ |
| 267 | source:Avoid redundant transpose of fused weights when weight_loader is called iteratively |
| 268 | """ |
| 269 | if expert_id is None and shard_id is None: |
| 270 | # MoE experts has been fused in disk |
| 271 | self._load_fused_experts_weight(param, loaded_weight) |
| 272 | return |
| 273 | if hasattr(param, "SHARD_ID_TO_SHARDED_DIM"): |
| 274 | SHARD_ID_TO_SHARDED_DIM = param.SHARD_ID_TO_SHARDED_DIM |
| 275 | elif current_platform.is_cuda() or current_platform.is_iluvatar() or current_platform.is_maca(): |
| 276 | SHARD_ID_TO_SHARDED_DIM = {"gate": 1, "down": 0, "up": 1} |
| 277 | else: |
| 278 | SHARD_ID_TO_SHARDED_DIM = {"gate": 0, "down": 1, "up": 0} |
| 279 | |
| 280 | if not (expert_id - self.expert_id_offset >= 0 and expert_id - self.expert_id_offset < self.num_local_experts): |
| 281 | return |
| 282 | if not param._is_initialized(): |
| 283 | param.initialize() |
| 284 | weight_need_transpose = getattr(param, "weight_need_transpose", False) |
| 285 | |
| 286 | if self.ep_size > 1 or weight_need_transpose: |
| 287 | loaded_weight = get_tensor(loaded_weight) |
| 288 | |
| 289 | if shard_id is None: |
| 290 | # 1.gate up fused in disk |
| 291 | if weight_need_transpose: |
| 292 | loaded_weight = loaded_weight.transpose([1, 0]) |
| 293 | output_size = param[expert_id - self.expert_id_offset].shape[SHARD_ID_TO_SHARDED_DIM["gate"]] |
| 294 | shard_offsets = [ |
| 295 | # (shard_id, shard_offset, shard_size) |
| 296 | ("gate", 0, output_size // 2 * self.tp_size), |
| 297 | ("up", output_size // 2 * self.tp_size, output_size // 2 * self.tp_size), |
| 298 | ] |
| 299 | |
| 300 | for shard_id, shard_offset, shard_size in shard_offsets: |
| 301 | loaded_weight_shard = slice_fn( |
| 302 | loaded_weight, SHARD_ID_TO_SHARDED_DIM[shard_id], shard_offset, shard_offset + shard_size |
| 303 | ) |
| 304 | self.weight_loader(param, loaded_weight_shard, expert_id, shard_id, "fused") |
| 305 | else: |
| 306 | if weight_need_transpose and source != "fused": |
| 307 | loaded_weight = loaded_weight.transpose([1, 0]) |
| 308 | # 2.gate up splited in disk |
| 309 | assert shard_id in ["gate", "down", "up"] |
| 310 | self._load_expert_weight( |
| 311 | param=param, |
| 312 | expert_id=expert_id, |
| 313 | loaded_weight=loaded_weight, |
| 314 | shard_id=shard_id, |
| 315 | shard_dim=SHARD_ID_TO_SHARDED_DIM[shard_id], |