Diffusion reasoningモデル「Dream 7B」を試す
kun432
🚀Dream 7B（拡散Reasoningモデル）を発表できることを嬉しく思います：現在までに開発された中で最も強力なオープン拡散大規模言語モデルです。
つまり、Dream 7Bは、
- 既存の拡散言語モデルを大幅に上回る性能を発揮し、
- 一般、数学、コーディング能力において、同程度の規模のトップクラスのAR言語モデルと同等以上の性能を発揮し、
- 優れた計画能力と柔軟な推論能力を示しています。
驚くべきことに、カウントダウンや数独のような計画的なタスクでは、タスク固有のトレーニングなしでも、DreamはQwen2.5 7BやLLaMA3 8Bを大幅に上回り、時には最新のDeepSeek V3さえも上回ります。
また、Dreamでは、出力を任意の順序で合成することも可能です。
正確な終了文を使った補完例
拡散タイムステップを調整することで、Dreamのパフォーマンスは、速度または品質のいずれかを柔軟に調整することができます。
kun432公式ブログ
hkunlp.github.io/blog/2025/dream/
GitHubレポジトリ
モデル
あと、HuggingFace Spaceにデモがある
kun432Colaboratory L4で、Dream-org/Dream-v0-Instruct-7Bを試す。
パッケージインストール
!pip install "transformers==4.46.2"
!pip install "torch==2.5.1" "torchvision==0.20.1" "torchaudio==2.5.1" --index-url https://download.pytorch.org/whl/cu124
モデルとトークナイザーのロード
import torch
from transformers import AutoModel, AutoTokenizer
model_path = "Dream-org/Dream-v0-Instruct-7B"
model = AutoModel.from_pretrained(
model_path,
torch_dtype=torch.bfloat16,
trust_remote_code=True
)
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
model = model.to("cuda").eval()
この時点でnvidia-smiを見るとVRAM消費は14.9GBぐらい。
Wed Apr 9 03:23:21 2025
+-----------------------------------------------------------------------------------------+
| NVIDIA-SMI 550.54.15 Driver Version: 550.54.15 CUDA Version: 12.4 |
|-----------------------------------------+------------------------+----------------------+
| GPU Name Persistence-M | Bus-Id Disp.A | Volatile Uncorr. ECC |
| Fan Temp Perf Pwr:Usage/Cap | Memory-Usage | GPU-Util Compute M. |
| | | MIG M. |
|=========================================+========================+======================|
| 0 NVIDIA L4 Off | 00000000:00:03.0 Off | 0 |
| N/A 49C P0 29W / 72W | 14887MiB / 23034MiB | 33% Default |
| | | N/A |
+-----------------------------------------+------------------------+----------------------+
では推論。GitHubのサンプルコードのまま。プロンプトだけ日本語にした。
messages = [
{
"role": "user",
"content": "Pythonのクラスを作成し、PyTorchのトレーナーを実装して、おもちゃのデータセットでモデルをトレーニングできるようにしてください。"
}
]
inputs = tokenizer.apply_chat_template(
messages, return_tensors="pt", return_dict=True, add_generation_prompt=True
)
input_ids = inputs.input_ids.to(device="cuda")
attention_mask = inputs.attention_mask.to(device="cuda")
output = model.diffusion_generate(
input_ids,
attention_mask=attention_mask,
max_new_tokens=512,
output_history=True,
return_dict_in_generate=True,
steps=512,
temperature=0.4,
top_p=0.95,
alg="entropy",
alg_temp=0.,
)
generations = [
tokenizer.decode(g[len(p) :].tolist())
for p, g in zip(input_ids, output.sequences)
]
print(generations[0].split(tokenizer.eos_token)[0])
結果。だいたい1分ぐらいかかった。なお、このコードではデモのような拡散的な様子は確認できない。
以下はPyTorchのトレーナーを実装して、おもちゃのデータセットでモデルをトレーニングするPythonのクラスの例です。
import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import Dataset, DataLoader class ToyDataset(Dataset): def __init__(self, data): self.data = data def __len__(self): return len(self.data) def __getitem__(self, index): return self.data[index] class ToyTrainer: def __init__(self, model, dataset, optimizer): self.model = model self.dataset = dataset self.optimizer = optimizer def train(self, epochs): for epoch in range(epochs): for data in self.dataset: self.model.train() self.optimizer.zero_grad() output = self.model(data) loss = self.loss(output, data) loss.backward() self.optimizer.step() class ToyModel(nn.Module): def __init__(self, data): super(ToyModel, self).__init__() self.data = data def forward(self, data): return self.data class ToyLoss(nn.Module): def __init__(self): super(ToyLoss, self).__init__() def forward(self, output, target): return torch.mean(output - target) data = torch.arange(5) dataset = ToyDataset(data) model = ToyModel(data) optimizer = optim.SGD(model.parameters(), lr=0.01) trainer = ToyTrainer(model, dataset, optimizer) trainer.train(10)このコードでは、
ToyDatasetはPyTorchのDatasetを相承して、おもちゃのデータセットを提供します。ToyTrainerは、モデルと、データセットと、オプトミライザを引いて、モデルをトレーニングするtrain方法を提供します。ToyModelは、PyTorchのnn.Moduleを相承して、おもちゃのモデルを提供します。ToyLossは、PyTorchのnn.Moduleを相承して、おもちゃのロス関数を提供します。最後に、おもちゃのデータセットと、モデルと、オプトミライザを生成し、ToyTrainerを生成し、モデルをトレーニングします。
推論後のVRAM消費は約17.6GB
Wed Apr 9 03:27:21 2025
+-----------------------------------------------------------------------------------------+
| NVIDIA-SMI 550.54.15 Driver Version: 550.54.15 CUDA Version: 12.4 |
|-----------------------------------------+------------------------+----------------------+
| GPU Name Persistence-M | Bus-Id Disp.A | Volatile Uncorr. ECC |
| Fan Temp Perf Pwr:Usage/Cap | Memory-Usage | GPU-Util Compute M. |
| | | MIG M. |
|=========================================+========================+======================|
| 0 NVIDIA L4 Off | 00000000:00:03.0 Off | 0 |
| N/A 71C P0 31W / 72W | 17627MiB / 23034MiB | 0% Default |
| | | N/A |
+-----------------------------------------+------------------------+----------------------+
デモっぽい感じが見たいなと思って、デモのコードをClaude-3.7-Sonnetに読ませて書いてもらった。
import time
import sys
from IPython.display import clear_output
# メッセージの設定
messages = [
{
"role": "user",
"content": "与えられた数値nからフィボナッチ数列のn番目の数値を返すPythonコードを書いてください。"
}
]
# テンプレートの適用
inputs = tokenizer.apply_chat_template(
messages, return_tensors="pt", return_dict=True, add_generation_prompt=True
)
input_ids = inputs.input_ids.to(device="cuda")
attention_mask = inputs.attention_mask.to(device="cuda")
# パラメータの設定
max_new_tokens = 512
steps = 512
temperature = 0.4
top_p = 0.95
alg = "origin" # origin / maskgit_plus / topk_margin / entropy から選択
alg_temp = 0.
# ステップカウンター（グローバル変数）
step_counter = 0
total_steps = steps + 1 # +1 for initial state
# 前のステップの状態を保存
previous_tokens = None
current_text = ""
# 拡散生成プロセスを追跡するためのカスタムフック関数
def token_hook(_, tokens_list, __):
"""
フック関数 - 引数は (None, tokens_list, None) の形式で渡されます
"""
global step_counter, previous_tokens, current_text
# 入力部分を除いたトークン列を取得
prompt_length = input_ids.shape[1]
generated_tokens = tokens_list[0][prompt_length:]
# 5ステップごとに表示を更新
if step_counter % 5 == 0 or step_counter == total_steps - 1:
# 画面をクリア（Jupyter/Colab環境）
clear_output(wait=True)
# マスクトークンの数をカウント
mask_count = (generated_tokens == tokenizer.mask_token_id).sum().item()
unmask_count = len(generated_tokens) - mask_count
# 生成されたテキストをデコード
current_text = tokenizer.decode(generated_tokens, skip_special_tokens=True)
# ヘッダー情報表示
print(f"===== 拡散進行状況: Step {step_counter}/{total_steps-1} ({step_counter/(total_steps-1)*100:.1f}%) =====")
print(f"マスク解除: {unmask_count}/{len(generated_tokens)} トークン解除済み")
# 前回との差分を確認して新しく生成されたトークンを把握
if previous_tokens is not None:
newly_unmasked = 0
for i, (prev, curr) in enumerate(zip(previous_tokens, generated_tokens)):
if prev == tokenizer.mask_token_id and curr != tokenizer.mask_token_id:
newly_unmasked += 1
#if newly_unmasked > 0:
# print(f"このステップで新たに {newly_unmasked} トークンのマスクが解除されました")
# 現在のテキストを表示
print("\n--- 現在の生成テキスト ---")
print(current_text)
print("-------------------------")
# 前のトークンを保存
previous_tokens = generated_tokens.clone()
# カウンターを増やす
step_counter += 1
return tokens_list
# 生成実行
print("生成開始...")
start_time = time.time()
output = model.diffusion_generate(
input_ids,
attention_mask=attention_mask,
max_new_tokens=max_new_tokens,
output_history=True,
return_dict_in_generate=True,
steps=steps,
temperature=temperature,
top_p=top_p,
alg=alg,
alg_temp=alg_temp,
generation_tokens_hook_func=token_hook # フックを設定
)
# 完了時間を記録
end_time = time.time()
print(f"\n生成完了！ 所要時間: {end_time - start_time:.2f}秒")
# 最終結果
prompt_length = input_ids.shape[1]
final_text = tokenizer.decode(output.sequences[0][prompt_length:], skip_special_tokens=True)
print("\n最終生成テキスト:")
print("=" * 80)
print(final_text)
print("=" * 80)
実行するとこんな感じになる。20倍速なので注意。実際はかなり遅い感はある。
kun432日本語もイケてるように思えるが、Xのポストなどにもある通り、数学的な問題やコード生成とかに向いているようで、普通の文章だと結構破綻するような印象がある。
(snip)
messages = [
{
"role": "user",
"content": "競馬の魅力について5つリストアップしてください。"
}
]
(snip)
結果
競馬の魅力は以下の5つです。最初に、競馬は競馬の速度と力のものです。また、競馬は心性です。次に、競馬の性質で、その性質を理解するものです。次に、競馬は機機性で、関心を引くものです。最後に、競馬は直感的性で、心動するものです。この5点は、競馬の魅力を示します。
まだこれは文章になってるだけマシかな、完全に壊れた生成結果になることも多い印象。
kun432個人的には拡散言語モデルのメリットがあまりわかっていない。
以前に見かけた拡散言語モデルはトータルの生成時間が自己回帰モデルよりも速いというものだった。
ただ、自己回帰モデルの場合、トータルの生成時間が長くても、ストリーミングなどのテクニックで見かけ上の速度が早くなっていればあんまり気にならないし、そもそも文章は頭から順に読んでいくので困ることもないんだよね。
拡散言語モデルだとちょっとそのあたりのテクニックが使えないのではないかと思うので、バッチ推論とかならいいのかなー、ぐらいのイメージしかない。
kun432それほどたくさんの試行回数ではないのだけども、一発目の推論のあとは、VRAM消費が安定してるように見える。
kun432書き忘れてたけと、そもそもそんなに速くないのでこんなもんなのかな？と思ってissue見てみたら上がってた
Hi there, the speed is related to max_new_tokens and steps. I just ran a test on one H800 GPU, and it costs 16s when setting max_new_tokens=512 and steps=512. So, I guess your speed seems reasonable considering the hardware difference.
そんなもんなのかな？