`OPTForCausalLM` 's `prefix` should not be `'model'` in `opt_modelling.py`

[sadra-barikbin]

Hi there!

OPTForCausalLM doesn't give prefix to OPTModel in opt_modelling.py while it's a positional argument to it.

text-generation-inference/server/text_generation_server/models/custom_modeling/opt_modeling.py

Lines 751 to 764 in 05c094f

	class OPTForCausalLM(OPTPreTrainedModel):
	def __init__(self, prefix, config, weights):
	super().__init__(config)
	if not prefix:
	prefix = "model"
	else:
	prefix = f"{prefix}.model"
	self.model = OPTModel(config, weights)
	self.lm_head = SpeculativeHead.load(
	config, prefix=f"{prefix}.decoder.embed_tokens", weights=weights
	)

text-generation-inference/server/text_generation_server/models/custom_modeling/opt_modeling.py

Lines 694 to 698 in 05c094f

	class OPTModel(OPTPreTrainedModel):
	def __init__(self, prefix: str, config: OPTConfig, weights):
	super().__init__(config)
	self.decoder = OPTDecoder(prefix, config, weights)
	# Initialize weights and apply final processing

[sadra-barikbin]

Besides, in v2.1.1, OPT's embedding is loaded using prefix

text-generation-inference/server/text_generation_server/models/custom_modeling/opt_modeling.py

Lines 440 to 442 in 4dfdb48

	self.embed_tokens = TensorParallelEmbedding(
	prefix="model.decoder.embed_tokens", weights=weights
	)

which raises this error:

 File "/opt/conda/lib/python3.10/site-packages/text_generation_server/models/opt.py", line 62, in __init__
    model = OPTForCausalLM(config, weights)

  File "/opt/conda/lib/python3.10/site-packages/text_generation_server/models/custom_modeling/opt_modeling.py", line 749, in __init__
    self.model = OPTModel(config, weights)

  File "/opt/conda/lib/python3.10/site-packages/text_generation_server/models/custom_modeling/opt_modeling.py", line 691, in __init__
    self.decoder = OPTDecoder(config, weights)

  File "/opt/conda/lib/python3.10/site-packages/text_generation_server/models/custom_modeling/opt_modeling.py", line 440, in __init__
    self.embed_tokens = TensorParallelEmbedding(

  File "/opt/conda/lib/python3.10/site-packages/text_generation_server/layers/tensor_parallel.py", line 230, in __init__
    weight = weights.get_partial_sharded(f"{prefix}.weight", dim=0)

  File "/opt/conda/lib/python3.10/site-packages/text_generation_server/utils/weights.py", line 89, in get_partial_sharded
    filename, tensor_name = self.get_filename(tensor_name)

  File "/opt/conda/lib/python3.10/site-packages/text_generation_server/utils/weights.py", line 64, in get_filename
    raise RuntimeError(f"weight {tensor_name} does not exist")

RuntimeError: weight model.decoder.embed_tokens.weight does not exist

With the latest changes on the main, It seems still it would raise error as CausalLM sets prefix to "" in

text-generation-inference/server/text_generation_server/models/causal_lm.py

Lines 556 to 557 in 05c094f

	prefix = ""
	model = model_class(prefix, config, weights)

and prefix becomes "model.decoder.embed_tokens"

[github-actions]

This issue is stale because it has been open 30 days with no activity. Remove stale label or comment or this will be closed in 5 days.

[sadra-barikbin]

Hi. This issue still persists. The OPT models in the hub do not have a 'model' at the beginning of their weights' names. This causes error when loading OPT which sets the prefix to 'model'. @danieldk

text-generation-inference/server/text_generation_server/models/custom_modeling/opt_modeling.py

Lines 755 to 758 in 133015f

	if not prefix:
	prefix = "model"
	else:
	prefix = f"{prefix}.model"

[sadra-barikbin]

Fixed by #2371