search

huggingface / text-generation-inference

Large Language Model Text Generation Inference
http://hf.co/docs/text-generation-inference
Apache License 2.0
8.75k stars 1.02k forks source link

Problem loading model with GPTNeoX architecture (weight gpt_neox.layers.0.attention.rotary_emb.inv_freq does not exist) #1460

Closed ZQ-Dev8 closed 7 months ago

ZQ-Dev8 commented 7 months ago

System Info

Information

Tasks

Reproduction

N/A, I am attempting to load the model via standard usage of TGI. The problem is related to the model, which uses the GPTNeoX architecture (pasted below). Note the model was not trained with flash attention, though the provided stacktrace seems to imply it needs to be (e.g. multiple instances of FlashNeoXAttention())

GPTNeoXForCausalLM(
  (gpt_neox): GPTNeoXModel(
    (embed_in): Embedding(50432, 2560)
    (emb_dropout): Dropout(p=0.0, inplace=False)
    (layers): ModuleList(
      (0-31): 32 x GPTNeoXLayer(
        (input_layernorm): LayerNorm((2560,), eps=1e-05, elementwise_affine=True)
        (post_attention_layernorm): LayerNorm((2560,), eps=1e-05, elementwise_affine=True)
        (post_attention_dropout): Dropout(p=0.0, inplace=False)
        (post_mlp_dropout): Dropout(p=0.0, inplace=False)
        (attention): GPTNeoXAttention(
          (rotary_emb): GPTNeoXRotaryEmbedding()
          (query_key_value): Linear(in_features=2560, out_features=7680, bias=True)
          (dense): Linear(in_features=2560, out_features=2560, bias=True)
          (attention_dropout): Dropout(p=0.0, inplace=False)
        )
        (mlp): GPTNeoXMLP(
          (dense_h_to_4h): Linear(in_features=2560, out_features=10240, bias=True)
          (dense_4h_to_h): Linear(in_features=10240, out_features=2560, bias=True)
          (act): GELUActivation()
        )
      )
    )
    (final_layer_norm): LayerNorm((2560,), eps=1e-05, elementwise_affine=True)
  )
  (embed_out): Linear(in_features=2560, out_features=50432, bias=False)

Expected behavior

I would expect the TGI server to start as with any other supported model. However, for this model I get the following error:

Container started successfully!
Running the text-generation-launcher command from /data directory inside the container...only local files will be used
2024-01-19T21:21:00.841715Z  INFO text_generation_launcher: Args { model_id: "/data/traclm-v1-3b-instruct", revision: None, validation_workers: 2, sharded: None, num_shard: Some(1), quantize: None, dtype: None, trust_remote_code: false, max_concurrent_requests: 128, max_best_of: 2, max_stop_sequences: 4, max_top_n_tokens: 5, max_input_length: 4096, max_total_tokens: 8192, waiting_served_ratio: 1.2, max_batch_prefill_tokens: 4096, max_batch_total_tokens: None, max_waiting_tokens: 20, hostname: "790d41f93e80", port: 80, shard_uds_path: "/tmp/text-generation-server", master_addr: "localhost", master_port: 29500, huggingface_hub_cache: Some("/data"), weights_cache_override: None, disable_custom_kernels: false, cuda_memory_fraction: 1.0, rope_scaling: None, rope_factor: None, json_output: false, otlp_endpoint: None, cors_allow_origin: [], watermark_gamma: None, watermark_delta: None, ngrok: false, ngrok_authtoken: None, ngrok_edge: None, env: false }
2024-01-19T21:21:00.841808Z  INFO download: text_generation_launcher: Starting download process.
2024-01-19T21:21:03.300505Z  INFO text_generation_launcher: Files are already present on the host. Skipping download.

2024-01-19T21:21:03.846038Z  INFO download: text_generation_launcher: Successfully downloaded weights.
2024-01-19T21:21:03.846486Z  INFO shard-manager: text_generation_launcher: Starting shard rank=0
2024-01-19T21:21:06.578815Z ERROR text_generation_launcher: Error when initializing model
Traceback (most recent call last):
  File "/opt/conda/bin/text-generation-server", line 8, in <module>
    sys.exit(app())
  File "/opt/conda/lib/python3.9/site-packages/typer/main.py", line 311, in __call__
    return get_command(self)(*args, **kwargs)
  File "/opt/conda/lib/python3.9/site-packages/click/core.py", line 1157, in __call__
    return self.main(*args, **kwargs)
  File "/opt/conda/lib/python3.9/site-packages/typer/core.py", line 778, in main
    return _main(
  File "/opt/conda/lib/python3.9/site-packages/typer/core.py", line 216, in _main
    rv = self.invoke(ctx)
  File "/opt/conda/lib/python3.9/site-packages/click/core.py", line 1688, in invoke
    return _process_result(sub_ctx.command.invoke(sub_ctx))
  File "/opt/conda/lib/python3.9/site-packages/click/core.py", line 1434, in invoke
    return ctx.invoke(self.callback, **ctx.params)
  File "/opt/conda/lib/python3.9/site-packages/click/core.py", line 783, in invoke
    return __callback(*args, **kwargs)
  File "/opt/conda/lib/python3.9/site-packages/typer/main.py", line 683, in wrapper
    return callback(**use_params)  # type: ignore
  File "/opt/conda/lib/python3.9/site-packages/text_generation_server/cli.py", line 83, in serve
    server.serve(
  File "/opt/conda/lib/python3.9/site-packages/text_generation_server/server.py", line 207, in serve
    asyncio.run(
  File "/opt/conda/lib/python3.9/asyncio/runners.py", line 44, in run
    return loop.run_until_complete(main)
  File "/opt/conda/lib/python3.9/asyncio/base_events.py", line 634, in run_until_complete
    self.run_forever()
  File "/opt/conda/lib/python3.9/asyncio/base_events.py", line 601, in run_forever
    self._run_once()
  File "/opt/conda/lib/python3.9/asyncio/base_events.py", line 1905, in _run_once
    handle._run()
  File "/opt/conda/lib/python3.9/asyncio/events.py", line 80, in _run
    self._context.run(self._callback, *self._args)
> File "/opt/conda/lib/python3.9/site-packages/text_generation_server/server.py", line 159, in serve_inner
    model = get_model(
  File "/opt/conda/lib/python3.9/site-packages/text_generation_server/models/__init__.py", line 175, in get_model
    return FlashNeoXSharded(
  File "/opt/conda/lib/python3.9/site-packages/text_generation_server/models/flash_neox.py", line 58, in __init__
    model = FlashGPTNeoXForCausalLM(config, weights)
  File "/opt/conda/lib/python3.9/site-packages/text_generation_server/models/custom_modeling/flash_neox_modeling.py", line 366, in __init__
    self.gpt_neox = FlashGPTNeoXModel(config, weights)
  File "/opt/conda/lib/python3.9/site-packages/text_generation_server/models/custom_modeling/flash_neox_modeling.py", line 308, in __init__
    [
  File "/opt/conda/lib/python3.9/site-packages/text_generation_server/models/custom_modeling/flash_neox_modeling.py", line 309, in <listcomp>
    FlashNeoXLayer(layer_id, config, weights)
  File "/opt/conda/lib/python3.9/site-packages/text_generation_server/models/custom_modeling/flash_neox_modeling.py", line 223, in __init__
    self.attention = FlashNeoxAttention(
  File "/opt/conda/lib/python3.9/site-packages/text_generation_server/models/custom_modeling/flash_neox_modeling.py", line 101, in __init__
    self.rotary_emb = PositionRotaryEmbedding.load(
  File "/opt/conda/lib/python3.9/site-packages/text_generation_server/utils/layers.py", line 625, in load
    inv_freq = weights.get_tensor(f"{prefix}.inv_freq")
  File "/opt/conda/lib/python3.9/site-packages/text_generation_server/utils/weights.py", line 75, in get_tensor
    filename, tensor_name = self.get_filename(tensor_name)
  File "/opt/conda/lib/python3.9/site-packages/text_generation_server/utils/weights.py", line 63, in get_filename
    raise RuntimeError(f"weight {tensor_name} does not exist")
RuntimeError: weight gpt_neox.layers.0.attention.rotary_emb.inv_freq does not exist

2024-01-19T21:21:07.151755Z ERROR shard-manager: text_generation_launcher: Shard complete standard error output:

Traceback (most recent call last):

  File "/opt/conda/bin/text-generation-server", line 8, in <module>
    sys.exit(app())

  File "/opt/conda/lib/python3.9/site-packages/text_generation_server/cli.py", line 83, in serve
    server.serve(

  File "/opt/conda/lib/python3.9/site-packages/text_generation_server/server.py", line 207, in serve
    asyncio.run(

  File "/opt/conda/lib/python3.9/asyncio/runners.py", line 44, in run
    return loop.run_until_complete(main)

  File "/opt/conda/lib/python3.9/asyncio/base_events.py", line 647, in run_until_complete
    return future.result()

  File "/opt/conda/lib/python3.9/site-packages/text_generation_server/server.py", line 159, in serve_inner
    model = get_model(

  File "/opt/conda/lib/python3.9/site-packages/text_generation_server/models/__init__.py", line 175, in get_model
    return FlashNeoXSharded(

  File "/opt/conda/lib/python3.9/site-packages/text_generation_server/models/flash_neox.py", line 58, in __init__
    model = FlashGPTNeoXForCausalLM(config, weights)

  File "/opt/conda/lib/python3.9/site-packages/text_generation_server/models/custom_modeling/flash_neox_modeling.py", line 366, in __init__
    self.gpt_neox = FlashGPTNeoXModel(config, weights)

  File "/opt/conda/lib/python3.9/site-packages/text_generation_server/models/custom_modeling/flash_neox_modeling.py", line 308, in __init__
    [

  File "/opt/conda/lib/python3.9/site-packages/text_generation_server/models/custom_modeling/flash_neox_modeling.py", line 309, in <listcomp>
    FlashNeoXLayer(layer_id, config, weights)

  File "/opt/conda/lib/python3.9/site-packages/text_generation_server/models/custom_modeling/flash_neox_modeling.py", line 223, in __init__
    self.attention = FlashNeoxAttention(

  File "/opt/conda/lib/python3.9/site-packages/text_generation_server/models/custom_modeling/flash_neox_modeling.py", line 101, in __init__
    self.rotary_emb = PositionRotaryEmbedding.load(

  File "/opt/conda/lib/python3.9/site-packages/text_generation_server/utils/layers.py", line 625, in load
    inv_freq = weights.get_tensor(f"{prefix}.inv_freq")

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

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

RuntimeError: weight gpt_neox.layers.0.attention.rotary_emb.inv_freq does not exist
 rank=0
2024-01-19T21:21:07.250391Z ERROR text_generation_launcher: Shard 0 failed to start
2024-01-19T21:21:07.250429Z  INFO text_generation_launcher: Shutting down shards
Error: ShardCannotStart
Failed to run text-generation-launcher.
ZQ-Dev8 commented 7 months ago

For additional context - I retrained the model, this time with flash attention, and the error remains the same:

RuntimeError: weight gpt_neox.layers.0.attention.rotary_emb.inv_freq does not exist
 rank=0
2024-01-20T14:17:26.514770Z ERROR text_generation_launcher: Shard 0 failed to start
2024-01-20T14:17:26.514815Z  INFO text_generation_launcher: Shutting down shards
Error: ShardCannotStart
Failed to run text-generation-launcher.
dwyatte commented 7 months ago

@Narsil I think this is probably the same as https://github.com/huggingface/text-generation-inference/issues/790 for Llama which you fixed by removing loading of position embeddings from weights in https://github.com/huggingface/text-generation-inference/pull/793. This fix seems simple enough but it looks like there was some concern about it breaking previous versions of the model / integration tests (but maybe everything was fine)

I think this probably broke when position embeddings were removed from the GPTNeoX weights (link to line diff from @ArthurZucker), which looks to have been released in transformers 4.35 so if you don't need Flash Attention, you might be able to downgrade transformers and re-export the model to get things working until fixed

ZQ-Dev8 commented 7 months ago

TYVM for the response. I can see the PR is held up - I'll attempt the downgrade/re-export you mentioned in the meantime.