Engine side fix for loading llama checkpoint fine-tuned with zero3

[minjiaz]

The loading of the Llama checkpoint is broken if we fine-tune it with zero-3. This issue prevents the loading of the Llama fine-tuned actor and critic models in DS-Chat for PPO stage 3 fine-tuning. The problem has been resolved by making the Llama container policy aware of zero-3.

[XuehaiPan]

Hi, the DeepSpeed Team, I encountered the same issue while enabling HybridEngine for the LLaMA model with ZeRO-3.

After several hours of debugging, I found a similar fix:

    def get_hidden_heads(self):
-       return self.client_module.self_attn.q_proj.weight.shape[1], \
+       return self.client_module.self_attn.q_proj.in_features, \
                self.client_module.self_attn.num_heads, \
                self.client_module.input_layernorm.variance_epsilon, \
-               self.client_module.mlp.gate_proj.weight.shape[0]
+               self.client_module.mlp.gate_proj.out_features

[XuehaiPan]

It would be nicer with temp variable + getattr.

Suggested change

	try: # for zero stage 3
	return self.client_module.self_attn.q_proj.weight.ds_shape[1], \
	self.client_module.self_attn.num_heads, \
	self.client_module.input_layernorm.variance_epsilon, \
	self.client_module.mlp.gate_proj.weight.ds_shape[0]
	except:
	return self.client_module.self_attn.q_proj.weight.shape[1], \
	self.client_module.self_attn.num_heads, \
	self.client_module.input_layernorm.variance_epsilon, \
	self.client_module.mlp.gate_proj.weight.shape[0]
	q_proj_weight = self.client_module.self_attn.q_proj.weight
	gate_proj_weight = self.client_module.mlp.gate_proj.weight
	return getattr(q_proj_weight, "ds_shape", q_proj_weight.shape)[1], \
	self.client_module.self_attn.num_heads, \
	self.client_module.input_layernorm.variance_epsilon, \
	getattr(gate_proj_weight, "ds_shape", gate_proj_weight.shape)[0]