The original issue mentioned potentially returning the log-likelihood of the tokens through some integration in the generate.py interface. I decided against implementing that support in this PR and instead returned the raw logits. My reasoning is that end users can then operate on these logits as needed.

Outputs from ./deepy.py generate.py ./configs/20B.yml ./configs/text_generation.yml on my branch

With return_logits: true

without

{
  "context": "",
  "text": "Entry DescriptionWild Agnes has released this extraordinary new",
  "length": 10,
  "finished": false,
  "message": null,
  "duration_seconds": 6.626634120941162
}

on main

{
  "context": "",
  "text": "Entry DescriptionWild Agnes has released this extraordinary new",
  "length": 10,
  "finished": false,
  "message": null,
  "duration_seconds": 6.523743629455566
}

I took a stab at writing some unit tests for this change but ran into an issue with mpu.get_model_parallel_src_rank() throwing an exception saying that I needed to configure init_process_group.