Open sunkun1997 opened 1 month ago
Same problem
This can (at least temporarily) be fixed by getting rid of the autocast at transformers/models/llama/modeling_llama.py And replacing everything from the # Force … comment in the forward pass to instead be:
freqs = (inv_freq_expanded @ position_ids_expanded).transpose(1, 2) emb = torch.cat((freqs, freqs), dim=-1) return emb.cos().to(dtype=x.dtype), emb.sin().to(dtype=x.dtype)
(this is basically taking the llama model back to commit 7628b3a0f40212c0f264233fc6da0d9c9cf88853 of the transformers package)
However, after doing this, there seems to still be a problem where the compiled (traced?, split?) model graph seems to not match the original:
[rank0]: Traceback (most recent call last):
[rank0]: File "/root/.local/share/code-server/extensions/ms-python.python-2022.4.1-universal/pythonFiles/lib/python/debugpy/_vendored/pydevd/pydevd.py", line 3385, in <module>
[rank0]: main()
[rank0]: File "/root/.local/share/code-server/extensions/ms-python.python-2022.4.1-universal/pythonFiles/lib/python/debugpy/_vendored/pydevd/pydevd.py", line 3378, in main
[rank0]: globals = debugger.run(setup['file'], None, None, is_module)
[rank0]: File "/root/.local/share/code-server/extensions/ms-python.python-2022.4.1-universal/pythonFiles/lib/python/debugpy/_vendored/pydevd/pydevd.py", line 2446, in run
[rank0]: return self._exec(is_module, entry_point_fn, module_name, file, globals, locals)
[rank0]: File "/root/.local/share/code-server/extensions/ms-python.python-2022.4.1-universal/pythonFiles/lib/python/debugpy/_vendored/pydevd/pydevd.py", line 2453, in _exec
[rank0]: pydev_imports.execfile(file, globals, locals) # execute the script
[rank0]: File "/root/.local/share/code-server/extensions/ms-python.python-2022.4.1-universal/pythonFiles/lib/python/debugpy/_vendored/pydevd/_pydev_bundle/_pydev_execfile.py", line 25, in execfile
[rank0]: exec(compile(contents + "\n", file, 'exec'), glob, loc)
[rank0]: File ".../pippy_llama.py", line 45, in <module>
[rank0]: pipe = pipeline(llama, example_args=(mb_inputs["input_ids"],), num_chunks=int(len(full_batch_prompts) / len(mb_inputs)))
[rank0]: File "/root/miniforge3/envs/sequin/lib/python3.10/site-packages/torch/distributed/pipelining/_IR.py", line 1187, in pipeline
[rank0]: return Pipe.from_tracing(
[rank0]: File "/root/miniforge3/envs/sequin/lib/python3.10/site-packages/torch/distributed/pipelining/_IR.py", line 1030, in from_tracing
[rank0]: pipe = Pipe._from_traced(
[rank0]: File "/root/miniforge3/envs/sequin/lib/python3.10/site-packages/torch/distributed/pipelining/_IR.py", line 734, in _from_traced
[rank0]: new_submod = _outline_submodules(submodule.graph)
[rank0]: File "/root/miniforge3/envs/sequin/lib/python3.10/site-packages/torch/distributed/pipelining/_unflatten.py", line 23, in _outline_submodules
[rank0]: ).run_outer()
[rank0]: File "/root/miniforge3/envs/sequin/lib/python3.10/site-packages/torch/export/unflatten.py", line 862, in run_outer
[rank0]: self.run_from(node_idx)
[rank0]: File "/root/miniforge3/envs/sequin/lib/python3.10/site-packages/torch/export/unflatten.py", line 942, in run_from
[rank0]: ).run_from(node_idx)
[rank0]: File "/root/miniforge3/envs/sequin/lib/python3.10/site-packages/torch/export/unflatten.py", line 942, in run_from
[rank0]: ).run_from(node_idx)
[rank0]: File "/root/miniforge3/envs/sequin/lib/python3.10/site-packages/torch/export/unflatten.py", line 942, in run_from
[rank0]: ).run_from(node_idx)
[rank0]: File "/root/miniforge3/envs/sequin/lib/python3.10/site-packages/torch/export/unflatten.py", line 919, in run_from
[rank0]: self.finalize_outputs()
[rank0]: File "/root/miniforge3/envs/sequin/lib/python3.10/site-packages/torch/export/unflatten.py", line 841, in finalize_outputs
[rank0]: _verify_graph_equivalence(self.cached_graph_module, self.module)
[rank0]: File "/root/miniforge3/envs/sequin/lib/python3.10/site-packages/torch/export/unflatten.py", line 567, in _verify_graph_equivalence
[rank0]: assert graph_dump(x.graph) == graph_dump(y.graph)
Was able to solve the above in my case by turning off the kv-cache in the model config. Perhaps this needs to be manually managed by the user outside of the traced module.
Can you tell me exactly which line you're replacing? And the way to turn off the kv-cache is revicing "use_cache": true into false in config.json?
Sorry the kv-cache thing was wrong. I was trying out gpt2 earlier to make sure I can at least run something.
Also had the wrong commit number earlier. Was talking about reverting this change in the transformer's library: https://github.com/huggingface/transformers/commit/d45f47ab7f7c31991bb98a0302ded59ab6adac31
Was able to resolve this by reverting transformers to the last December 2023 commit that passes all tests (3b7675b2b844b02d4821b827871a21ad16dd446c) and the PiPPy v0.2.0 tag. If you need batch chat template decoding then you need to go find the updated utils tokenization base file and the init.py file for that folder accordingly as well.
torch version: 2.5.0.dev20240616+cu121 python version: python 3.8
I run the llama example with torchrun --nproc-per-node 2 pippy_llama.py. It got an Error