Lightning-AI / Lightning-AI/lightning-thunder

Unknown attribute _base inside Megatron core

Open
#753 12 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

nemo program-coverage triage review
Dominant language
Python
Stars
1.5k
Forks
121
PR merge metrics
No merged PRs in 30d

Description

## 🚀 Model / language coverage

Running Megatron GPT from NeMo, we seem to have issues with [this line from Megatron core](https://github.com/NVIDIA/Megatron-LM/blob/c7a1f82d761577e6ca0338d3521eac82f2aa0904/megatron/core/utils.py#L147). [Some context from the caller in this particular case](https://github.com/NVIDIA/Megatron-LM/blob/c7a1f82d761577e6ca0338d3521eac82f2aa0904/megatron/core/transformer/transformer_block.py#L346) might be illuminating.

```
[rank0]: Traceback (most recent call last):
[rank0]: File "/home/tfogal/dev/nemo/examples/nlp/language_modeling/tuning/megatron_gpt_finetuning.py", line 88, in
[rank0]: main()
[rank0]: File "/home/tfogal/dev/nemo/nemo/core/config/hydra_runner.py", line 129, in wrapper
[rank0]: _run_hydra(
[rank0]: File "/home/tfogal/env/lib/python3.10/site-packages/hydra/_internal/utils.py", line 394, in _run_hydra
[rank0]: _run_app(
[rank0]: File "/home/tfogal/env/lib/python3.10/site-packages/hydra/_internal/utils.py", line 457, in _run_app
[rank0]: run_and_report(
[rank0]: File "/home/tfogal/env/lib/python3.10/site-packages/hydra/_internal/utils.py", line 223, in run_and_report
[rank0]: raise ex
[rank0]: File "/home/tfogal/env/lib/python3.10/site-packages/hydra/_internal/utils.py", line 220, in run_and_report
[rank0]: return func()
[rank0]: File "/home/tfogal/env/lib/python3.10/site-packages/hydra/_internal/utils.py", line 458, in
[rank0]: lambda: hydra.run(
[rank0]: File "/home/tfogal/env/lib/python3.10/site-packages/hydra/_internal/hydra.py", line 132, in run
[rank0]: _ = ret.return_value
[rank0]: File "/home/tfogal/env/lib/python3.10/site-packages/hydra/core/utils.py", line 260, in return_value
[rank0]: raise self._return_value
[rank0]: File "/home/tfogal/env/lib/python3.10/site-packages/hydra/core/utils.py", line 186, in run_job
[rank0]: ret.return_value = task_function(task_cfg)
[rank0]: File "/home/tfogal/dev/nemo/examples/nlp/language_modeling/tuning/megatron_gpt_finetuning.py", line 84, in main
[rank0]: trainer.fit(model)
[rank0]: File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/trainer/trainer.py", line 543, in fit
[rank0]: call._call_and_handle_interrupt(
[rank0]: File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/trainer/call.py", line 43, in _call_and_handle_interrupt
[rank0]: return trainer.strategy.launcher.launch(trainer_fn, *args, trainer=trainer, **kwargs)
[rank0]: File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/strategies/launchers/subprocess_script.py", line 105, in launch
[rank0]: return function(*args, **kwargs)
[rank0]: File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/trainer/trainer.py", line 579, in _fit_impl
[rank0]: self._run(model, ckpt_path=ckpt_path)
[rank0]: File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/trainer/trainer.py", line 986, in _run
[rank0]: results = self._run_stage()
[rank0]: File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/trainer/trainer.py", line 1028, in _run_stage
[rank0]: self._run_sanity_check()
[rank0]: File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/trainer/trainer.py", line 1057, in _run_sanity_check
[rank0]: val_loop.run()
[rank0]: File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/loops/utilities.py", line 182, in _decorator
[rank0]: return loop_run(self, *args, **kwargs)
[rank0]: File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/loops/evaluation_loop.py", line 135, in run
[rank0]: self._evaluation_step(batch, batch_idx, dataloader_idx, dataloader_iter)
[rank0]: File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/loops/evaluation_loop.py", line 396, in _evaluation_step
[rank0]: output = call._call_strategy_hook(trainer, hook_name, *step_args)
[rank0]: File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/trainer/call.py", line 311, in _call_strategy_hook
[rank0]: output = fn(*args, **kwargs)
[rank0]: File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/strategies/strategy.py", line 411, in validation_step
[rank0]: return self.lightning_module.validation_step(*args, **kwargs)
[rank0]: File "/home/tfogal/dev/nemo/nemo/collections/nlp/models/language_modeling/megatron_gpt_sft_model.py", line 434, in validation_step
[rank0]: return self.inference_step(dataloader_iter, 'validation')
[rank0]: File "/home/tfogal/dev/nemo/nemo/collections/nlp/models/language_modeling/megatron_gpt_sft_model.py", line 444, in inference_step
[rank0]: outputs = self.inference_step_validation_call(batch, batch_idx, data_cfg, dataloader_idx)
[rank0]: File "/home/tfogal/dev/nemo/nemo/collections/nlp/models/language_modeling/megatron_gpt_sft_model.py", line 464, in inference_step_validation_call
[rank0]: loss = super().validation_step(itertools.chain([batch]), dataloader_idx)
[rank0]: File "/home/tfogal/dev/nemo/nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py", line 1370, in validation_step
[rank0]: loss = self.fwd_bwd_step(dataloader_iter, True, first_val_step)
[rank0]: File "/home/tfogal/dev/nemo/nemo/collections/nlp/models/language_modeling/megatron_gpt_sft_model.py", line 381, in fwd_bwd_step
[rank0]: losses_reduced_per_micro_batch = fwd_bwd_function(
[rank0]: File "/home/tfogal/env/lib/python3.10/site-packages/megatron/core/pipeline_parallel/schedules.py", line 392, in forward_backward_no_pipelining
[rank0]: output_tensor, num_tokens = forward_step(
[rank0]: File "/home/tfogal/env/lib/python3.10/site-packages/megatron/core/pipeline_parallel/schedules.py", line 217, in forward_step
[rank0]: output_tensor, loss_func = forward_step_func(data_iterator, model)
[rank0]: File "/home/tfogal/dev/nemo/nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py", line 1232, in fwd_output_and_loss_func
[rank0]: output_tensor = model(**forward_args)
[rank0]: File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1716, in _wrapped_call_impl
[rank0]: return self._call_impl(*args, **kwargs)
[rank0]: File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1727, in _call_impl
[rank0]: return forward_call(*args, **kwargs)
[rank0]: File "/home/tfogal/dev/thunder/thunder/core/module.py", line 61, in forward
[rank0]: res = self._forward_fn(*args, **kwargs)
[rank0]: File "/home/tfogal/dev/thunder/thunder/__init__.py", line 669, in fn_
[rank0]: cache_entry, inps, pro_to_epi = get_computation_and_inputs(*args, **kwargs)
[rank0]: File "/home/tfogal/dev/thunder/thunder/__init__.py", line 223, in cache_info_wrapper
[rank0]: res = fn(*args, **kwargs)
[rank0]: File "/home/tfogal/dev/thunder/thunder/__init__.py", line 503, in get_computation_and_inputs
[rank0]: jit_results: TraceResults = interpreter(
[rank0]: File "/home/tfogal/dev/thunder/thunder/__init__.py", line 211, in _general_frontend
[rank0]: return thunder_general_jit(fn, args, kwargs, sharp_edges=sharp_edges, record_history=record_history)
[rank0]: File "/home/tfogal/dev/thunder/thunder/core/jit_ext.py", line 1743, in thunder_general_jit
[rank0]: result = jfn(*args, **kwargs)
[rank0]: File "/home/tfogal/dev/thunder/thunder/core/interpreter.py", line 6695, in fn_
[rank0]: raise e
[rank0]: File "/home/tfogal/dev/thunder/thunder/core/interpreter.py", line 6663, in fn_2
[rank0]: return fn(*args, **kwargs)
[rank0]: File "/home/tfogal/dev/thunder/thunder/core/interpreter.py", line 6060, in _impl
[rank0]: return fn.__func__(fn.__self__, *args, **kwargs)
[rank0]: File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1716, in _wrapped_call_impl
[rank0]: return self._call_impl(*args, **kwargs)
[rank0]: File "/home/tfogal/dev/thunder/thunder/core/interpreter.py", line 6060, in _impl
[rank0]: return fn.__func__(fn.__self__, *args, **kwargs)
[rank0]: File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1727, in _call_impl
[rank0]: return forward_call(*args, **kwargs)
[rank0]: File "/home/tfogal/dev/thunder/thunder/core/interpreter.py", line 6060, in _impl
[rank0]: return fn.__func__(fn.__self__, *args, **kwargs)
[rank0]: File "/home/tfogal/env/lib/python3.10/site-packages/megatron/core/models/gpt/gpt_model.py", line 190, in forward
[rank0]: hidden_states = self.decoder(
[rank0]: File "/home/tfogal/dev/thunder/thunder/core/interpreter.py", line 6060, in _impl
[rank0]: return fn.__func__(fn.__self__, *args, **kwargs)
[rank0]: File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1716, in _wrapped_call_impl
[rank0]: return self._call_impl(*args, **kwargs)
[rank0]: File "/home/tfogal/dev/thunder/thunder/core/interpreter.py", line 6060, in _impl
[rank0]: return fn.__func__(fn.__self__, *args, **kwargs)
[rank0]: File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1727, in _call_impl
[rank0]: return forward_call(*args, **kwargs)
[rank0]: File "/home/tfogal/dev/thunder/thunder/core/interpreter.py", line 6060, in _impl
[rank0]: return fn.__func__(fn.__self__, *args, **kwargs)
[rank0]: File "/home/tfogal/env/lib/python3.10/site-packages/megatron/core/transformer/transformer_block.py", line 361, in forward
[rank0]: hidden_states = make_viewless_tensor(
[rank0]: File "/home/tfogal/env/lib/python3.10/site-packages/megatron/core/utils.py", line 147, in make_viewless_tensor
[rank0]: if inp._base is None:
[rank0]: File "/home/tfogal/dev/thunder/thunder/core/interpreter.py", line 6060, in _impl
[rank0]: return fn.__func__(fn.__self__, *args, **kwargs)
[rank0]: File "/home/tfogal/dev/thunder/thunder/core/interpreter.py", line 1272, in wrapping_wrapper
[rank0]: res = ufn(*uargs, **ukwargs)
[rank0]: File "/home/tfogal/dev/thunder/thunder/core/jit_ext.py", line 704, in wrapper
[rank0]: return fn(*args, **kwargs)
[rank0]: File "/home/tfogal/dev/thunder/thunder/core/proxies.py", line 1312, in __getattr__
[rank0]: baseutils.check(method_or_value is not None, lambda: f"Unknown attribute {attr}", exception_type=AttributeError)
[rank0]: File "/home/tfogal/dev/thunder/thunder/core/baseutils.py", line 103, in check
[rank0]: raise exception_type(s())
[rank0]: AttributeError: Unknown attribute _base. Did you mean: '_return_value'?
```

### Pitch

This is coming from a finetuning case of the Megatron GPT network. To setup and run the network, see #344.

### Alternatives / Potential work-arounds

Marking triage review for help

### Minimal Repro

(help wanted)

cc @apaz-cli @tfogal

Contributor guide

No contributing guide indexed for this repository

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start with thunder/core/proxies.py, especially __getattr__, and trace the _base access from Megatron's megatron/core/utils.py:147 and transformer_block.py:346. Reproduce it through the finetuning entry point in examples/nlp/language_modeling/tuning/megatron_gpt_finetuning.py described in #344. Done means this Megatron GPT path no longer raises Unknown attribute _base.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
compilers, machine-learning
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.