Potential issue in prompt handling in `generate()` in `torchtune/recipes/generate.py`

[insop]

Potential issue in prompt handling in generate()

Description

generate() in torchtune/recipes/generate.py receives a cfg dict that includes 'prompt', 'system', 'user' and other keys. However, it only passes cfg.prompt to self.convert_prompt_to_tokens(). This causes a failure since self.convert_prompt_to_tokens looks for cfg.system and cfg.user keys.

Documentation

There are two documents that describe using generate():

• link1
• link2

Potential solution

A potential solution could be passing cfg to self.convert_prompt_to_tokens() instead of cfg.prompt only.

Error message

Here is an error message:

Traceback (most recent call last):
  File "/usr/local/bin/tune", line 8, in <module>
    sys.exit(main())
  File "/workspace/torchtune/torchtune/_cli/tune.py", line 49, in main
    parser.run(args)
  File "/workspace/torchtune/torchtune/_cli/tune.py", line 43, in run
    args.func(args)
  File "/workspace/torchtune/torchtune/_cli/run.py", line 214, in _run_cmd
    self._run_single_device(args, is_builtin=is_builtin)
  File "/workspace/torchtune/torchtune/_cli/run.py", line 108, in _run_single_device
    runpy.run_path(str(args.recipe), run_name="__main__")
  File "/usr/lib/python3.10/runpy.py", line 289, in run_path
    return _run_module_code(code, init_globals, run_name,
  File "/usr/lib/python3.10/runpy.py", line 96, in _run_module_code
    _run_code(code, mod_globals, init_globals,
  File "/usr/lib/python3.10/runpy.py", line 86, in _run_code
    exec(code, run_globals)
  File "/workspace/torchtune/recipes/generate.py", line 203, in <module>
  File "/workspace/torchtune/torchtune/config/_parse.py", line 99, in wrapper
    sys.exit(recipe_main(conf))
  File "/workspace/torchtune/recipes/generate.py", line 199, in main
    
  File "/usr/local/lib/python3.10/dist-packages/torch/utils/_contextlib.py", line 116, in decorate_context
    return func(*args, **kwargs)
  File "/workspace/torchtune/recipes/generate.py", line 123, in generate
    cfg.prompt,
  File "/workspace/torchtune/recipes/generate.py", line 109, in convert_prompt_to_tokens
    if "system" in prompt and prompt["system"] is not None:
TypeError: argument of type 'NoneType' is not iterable

[RdoubleA]

@insop Appreciate the detailed write up. Could you share the generation config you used that produced this error?

[insop]

Thank you @RdoubleA

Command used:

tune run generate --config custom_generation_config.yaml

Config (custom_generation_config.yaml) used:

model:
  _component_: torchtune.models.llama3_1.llama3_1_8b

checkpointer:
  _component_: torchtune.training.FullModelHFCheckpointer
  checkpoint_dir: PATH TO ...models--meta-llama--Llama-3.1-8B
  checkpoint_files: [
    model-00001-of-00004.safetensors,
    model-00002-of-00004.safetensors,
    model-00003-of-00004.safetensors,
    model-00004-of-00004.safetensors
  ]

  output_dir: PATH TO OUTPUT
  model_type: LLAMA3

# Make sure to update the tokenizer path to the right
# checkpoint directory as well
tokenizer:
  _component_: torchtune.models.llama3.llama3_tokenizer
  path: PATH TO TOKENIZER

# Environment
device: cuda
dtype: bf16
enable_activation_checkpointing: False  # True reduces memory
seed: 1234 # It is not recommended to change this seed, b/c it matches EleutherAI's default seed

# Generation arguments; defaults taken from gpt-fast
prompt:
system: null
user: "My name is ... "
max_new_tokens: 2048
temperature: 0.6 # 0.8 and 0.6 are popular values to try
top_k: 300

enable_kv_cache: True

# Quantization specific args
quantizer: null

[RdoubleA]

The indents are intentional in the example generation configs. It indicates to the config object that the prompt: field should be a dictionary. Currently in your config, the prompt field is being treated as empty or None

prompt:
system: null
user: "My name is ... "

This will create the following cfg object:

cfg["prompt"] = None
cfg["system"] = None
cfg["user"] = "My name is ..."

but what the generation recipe expects is that prompt is a nested dictionary with system and user fields. So your config should look like so:

prompt:
  system: null
  user: "My name is ..."

which will create the following object:

cfg["prompt"] = {"system": None, "user": "My name is ..."}

[insop]

Thank you, that's great to know.
So with the config update, it seems good.

However, there could be an issue.

This is an example from this document.

tune run generate --config `find . -name custom_generation_config.yaml` prompt="Hello, my name is"

If you do as suggested in the document, then it will fail with the following errors. It is because we have cfg['prompt'] = "Hello, my name is".

  File "/workspace/torchtune/recipes/generate.py", line 114, in convert_prompt_to_tokens
    Message(role="user", content=prompt["user"]),

(update) just one improvement would be goodl
Two improvements I can think of:

1. update the example of the document.

tune run generate --config ./custom_generation_config.yaml \
prompt.user="Hello, my name is"

[RdoubleA]

Ah, you're right. That's a good catch. It should be updated to prompt.user as you suggested.

You're welcome to put up a PR if you'd like, if not I can put up a quick fix. Let me know.

[insop]

Thank you @RdoubleA.

I can make a quick fix shortly.