Buckets:

EmergentTwilight
/

MMCoT

Files

xet

EmergentTwilight/MMCoT / OURS /tests /test_rl_distributed.py

EmergentTwilight

24 days ago

download

raw

8.38 kB

import tempfile
import unittest
from pathlib import Path
from types import SimpleNamespace

import torch
import torch.distributed as dist
import torch.multiprocessing as mp

from rl.distributed import assigned_group_indices, DistributedContext
from rl.config import RLConfig
from rl.engine import execute_actor_update
from rl.objective import TrajectoryPolicyInputs
from rl.policy import PolicyScore


class _Scorer:
    def __init__(self, actor):
        self.actor = actor

    def score(self, trajectory, *, advantage):
        current = self.actor.weight.reshape(1)
        if trajectory == "b":
            current = -current
        empty = current.new_empty((0, 0))
        return PolicyScore(
            inputs=TrajectoryPolicyInputs(
                advantage=advantage,
                text_current_log_probs=current,
                text_old_log_probs=torch.zeros(1),
                visual_current_factor_log_probs=empty,
                visual_old_factor_log_probs=empty,
            ),
            text_reference_log_probs=torch.zeros(1),
            visual_reference_factor_log_probs=empty,
            text_replay_log_probs=current.detach(),
            visual_replay_factor_log_probs=empty,
            behavior_log_prob_max_abs_error=0.0,
            text_behavior_log_prob_max_abs_error=0.0,
            visual_behavior_log_prob_max_abs_error=0.0,
            text_behavior_log_prob_mean_abs_error=0.0,
            text_behavior_log_prob_p99_abs_error=0.0,
            visual_behavior_log_prob_mean_abs_error=0.0,
            visual_behavior_log_prob_p99_abs_error=0.0,
            visual_behavior_log_prob_step_signed_mean_abs_p99=0.0,
            text_actions=1,
            visual_steps=0,
        )


def _group():
    return SimpleNamespace(
        trajectories=(
            SimpleNamespace(trajectory="a"),
            SimpleNamespace(trajectory="b"),
        ),
        advantages=torch.tensor([1.0, -1.0]),
    )


def _engine_worker(rank: int, rendezvous: str, output: str) -> None:
    dist.init_process_group(
        "gloo",
        init_method=f"file://{rendezvous}",
        rank=rank,
        world_size=2,
    )
    actor = torch.nn.Linear(1, 1, bias=False)
    actor.weight.data.zero_()
    optimizer = torch.optim.SGD(actor.parameters(), lr=0.1)
    scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lambda _step: 1.0)
    config = RLConfig()
    config.rollout.group_size = 2
    config.trainer.world_size = 2
    config.reference.beta = 0.0
    result = execute_actor_update(
        groups=[_group()],
        policy_scorer=_Scorer(actor),
        config=config,
        parameters=list(actor.parameters()),
        actor=actor,
        optimizer=optimizer,
        scheduler=scheduler,
        parallel=DistributedContext(rank=rank, local_rank=rank, world_size=2),
    )
    torch.save(
        {
            "weight": actor.weight.detach(),
            "grad_norm": result.grad_norm,
            "records": len(result.policy_records),
        },
        Path(output) / f"engine-rank-{rank}.pt",
    )
    dist.destroy_process_group()


def _stage_error_worker(rank: int, rendezvous: str, output: str) -> None:
    dist.init_process_group(
        "gloo",
        init_method=f"file://{rendezvous}",
        rank=rank,
        world_size=2,
    )
    context = DistributedContext(rank=rank, local_rank=rank, world_size=2)
    try:
        context.gather_stage(
            rank,
            "rank 1: RuntimeError: fixture" if rank == 1 else None,
            stage="rollout",
        )
    except RuntimeError as error:
        message = str(error)
    else:
        message = "missing error"
    (Path(output) / f"stage-rank-{rank}.txt").write_text(message)
    dist.destroy_process_group()


def _gradient_worker(rank: int, rendezvous: str, output: str) -> None:
    dist.init_process_group(
        "gloo",
        init_method=f"file://{rendezvous}",
        rank=rank,
        world_size=2,
    )
    shared = torch.nn.Parameter(torch.tensor([0.0]))
    rank_zero_only = torch.nn.Parameter(torch.tensor([0.0]))
    globally_unused = torch.nn.Parameter(torch.tensor([0.0]))
    shared.grad = torch.tensor([float(rank + 1)])
    if rank == 0:
        rank_zero_only.grad = torch.tensor([4.0])
    context = DistributedContext(rank=rank, local_rank=rank, world_size=2)
    context.synchronize_gradients((shared, rank_zero_only, globally_unused))
    torch.save(
        {
            "shared": shared.grad,
            "rank_zero_only": rank_zero_only.grad,
            "globally_unused": globally_unused.grad,
        },
        Path(output) / f"rank-{rank}.pt",
    )
    dist.destroy_process_group()


class RLDistributedTest(unittest.TestCase):
    def test_group_members_are_evenly_partitioned_without_overlap(self):
        contexts = [
            DistributedContext(rank=rank, local_rank=rank, world_size=4)
            for rank in range(4)
        ]
        shards = [assigned_group_indices(8, context) for context in contexts]

        self.assertEqual(shards, [(0, 4), (1, 5), (2, 6), (3, 7)])
        self.assertEqual(sorted(index for shard in shards for index in shard), list(range(8)))

        uneven = [
            assigned_group_indices(
                4,
                DistributedContext(rank=rank, local_rank=rank, world_size=3),
            )
            for rank in range(3)
        ]
        self.assertEqual(uneven, [(0, 3), (1,), (2,)])

    def test_gradient_sum_handles_rank_local_and_globally_unused_parameters(self):
        with tempfile.TemporaryDirectory() as directory:
            rendezvous = str(Path(directory) / "rendezvous")
            mp.spawn(
                _gradient_worker,
                args=(rendezvous, directory),
                nprocs=2,
                join=True,
            )
            results = [
                torch.load(
                    Path(directory) / f"rank-{rank}.pt",
                    weights_only=True,
                )
                for rank in range(2)
            ]

        for result in results:
            self.assertEqual(result["shared"].item(), 3.0)
            self.assertEqual(result["rank_zero_only"].item(), 4.0)
            self.assertIsNone(result["globally_unused"])

    def test_two_rank_actor_update_matches_single_worker_global_group(self):
        actor = torch.nn.Linear(1, 1, bias=False)
        actor.weight.data.zero_()
        optimizer = torch.optim.SGD(actor.parameters(), lr=0.1)
        scheduler = torch.optim.lr_scheduler.LambdaLR(
            optimizer, lambda _step: 1.0
        )
        config = RLConfig()
        config.rollout.group_size = 2
        config.reference.beta = 0.0
        expected = execute_actor_update(
            groups=[_group()],
            policy_scorer=_Scorer(actor),
            config=config,
            parameters=list(actor.parameters()),
            actor=actor,
            optimizer=optimizer,
            scheduler=scheduler,
        )

        with tempfile.TemporaryDirectory() as directory:
            mp.spawn(
                _engine_worker,
                args=(str(Path(directory) / "engine-rendezvous"), directory),
                nprocs=2,
                join=True,
            )
            results = [
                torch.load(
                    Path(directory) / f"engine-rank-{rank}.pt",
                    weights_only=True,
                )
                for rank in range(2)
            ]

        for result in results:
            self.assertTrue(torch.equal(result["weight"], actor.weight.detach()))
            self.assertEqual(result["grad_norm"], expected.grad_norm)
            self.assertEqual(result["records"], 2)

    def test_rank_local_stage_error_is_reported_by_every_worker(self):
        with tempfile.TemporaryDirectory() as directory:
            mp.spawn(
                _stage_error_worker,
                args=(str(Path(directory) / "stage-rendezvous"), directory),
                nprocs=2,
                join=True,
            )
            messages = [
                (Path(directory) / f"stage-rank-{rank}.txt").read_text()
                for rank in range(2)
            ]

        self.assertEqual(messages[0], messages[1])
        self.assertIn("distributed rollout failed", messages[0])
        self.assertIn("rank 1: RuntimeError: fixture", messages[0])


if __name__ == "__main__":
    unittest.main()

Xet Storage Details

Size:: 8.38 kB
Xet hash:: 6d12d2f2e21d965e16e85d9f59ec5aca3930c2a5a4a9b741d2ddbee59a209d95

Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.

	import tempfile
	import unittest
	from pathlib import Path
	from types import SimpleNamespace

	import torch
	import torch.distributed as dist
	import torch.multiprocessing as mp

	from rl.distributed import assigned_group_indices, DistributedContext
	from rl.config import RLConfig
	from rl.engine import execute_actor_update
	from rl.objective import TrajectoryPolicyInputs
	from rl.policy import PolicyScore


	class _Scorer:
	def __init__(self, actor):
	self.actor = actor

	def score(self, trajectory, *, advantage):
	current = self.actor.weight.reshape(1)
	if trajectory == "b":
	current = -current
	empty = current.new_empty((0, 0))
	return PolicyScore(
	inputs=TrajectoryPolicyInputs(
	advantage=advantage,
	text_current_log_probs=current,
	text_old_log_probs=torch.zeros(1),
	visual_current_factor_log_probs=empty,
	visual_old_factor_log_probs=empty,
	),
	text_reference_log_probs=torch.zeros(1),
	visual_reference_factor_log_probs=empty,
	text_replay_log_probs=current.detach(),
	visual_replay_factor_log_probs=empty,
	behavior_log_prob_max_abs_error=0.0,
	text_behavior_log_prob_max_abs_error=0.0,
	visual_behavior_log_prob_max_abs_error=0.0,
	text_behavior_log_prob_mean_abs_error=0.0,
	text_behavior_log_prob_p99_abs_error=0.0,
	visual_behavior_log_prob_mean_abs_error=0.0,
	visual_behavior_log_prob_p99_abs_error=0.0,
	visual_behavior_log_prob_step_signed_mean_abs_p99=0.0,
	text_actions=1,
	visual_steps=0,
	)


	def _group():
	return SimpleNamespace(
	trajectories=(
	SimpleNamespace(trajectory="a"),
	SimpleNamespace(trajectory="b"),
	),
	advantages=torch.tensor([1.0, -1.0]),
	)


	def _engine_worker(rank: int, rendezvous: str, output: str) -> None:
	dist.init_process_group(
	"gloo",
	init_method=f"file://{rendezvous}",
	rank=rank,
	world_size=2,
	)
	actor = torch.nn.Linear(1, 1, bias=False)
	actor.weight.data.zero_()
	optimizer = torch.optim.SGD(actor.parameters(), lr=0.1)
	scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lambda _step: 1.0)
	config = RLConfig()
	config.rollout.group_size = 2
	config.trainer.world_size = 2
	config.reference.beta = 0.0
	result = execute_actor_update(
	groups=[_group()],
	policy_scorer=_Scorer(actor),
	config=config,
	parameters=list(actor.parameters()),
	actor=actor,
	optimizer=optimizer,
	scheduler=scheduler,
	parallel=DistributedContext(rank=rank, local_rank=rank, world_size=2),
	)
	torch.save(
	{
	"weight": actor.weight.detach(),
	"grad_norm": result.grad_norm,
	"records": len(result.policy_records),
	},
	Path(output) / f"engine-rank-{rank}.pt",
	)
	dist.destroy_process_group()


	def _stage_error_worker(rank: int, rendezvous: str, output: str) -> None:
	dist.init_process_group(
	"gloo",
	init_method=f"file://{rendezvous}",
	rank=rank,
	world_size=2,
	)
	context = DistributedContext(rank=rank, local_rank=rank, world_size=2)
	try:
	context.gather_stage(
	rank,
	"rank 1: RuntimeError: fixture" if rank == 1 else None,
	stage="rollout",
	)
	except RuntimeError as error:
	message = str(error)
	else:
	message = "missing error"
	(Path(output) / f"stage-rank-{rank}.txt").write_text(message)
	dist.destroy_process_group()


	def _gradient_worker(rank: int, rendezvous: str, output: str) -> None:
	dist.init_process_group(
	"gloo",
	init_method=f"file://{rendezvous}",
	rank=rank,
	world_size=2,
	)
	shared = torch.nn.Parameter(torch.tensor([0.0]))
	rank_zero_only = torch.nn.Parameter(torch.tensor([0.0]))
	globally_unused = torch.nn.Parameter(torch.tensor([0.0]))
	shared.grad = torch.tensor([float(rank + 1)])
	if rank == 0:
	rank_zero_only.grad = torch.tensor([4.0])
	context = DistributedContext(rank=rank, local_rank=rank, world_size=2)
	context.synchronize_gradients((shared, rank_zero_only, globally_unused))
	torch.save(
	{
	"shared": shared.grad,
	"rank_zero_only": rank_zero_only.grad,
	"globally_unused": globally_unused.grad,
	},
	Path(output) / f"rank-{rank}.pt",
	)
	dist.destroy_process_group()


	class RLDistributedTest(unittest.TestCase):
	def test_group_members_are_evenly_partitioned_without_overlap(self):
	contexts = [
	DistributedContext(rank=rank, local_rank=rank, world_size=4)
	for rank in range(4)
	]
	shards = [assigned_group_indices(8, context) for context in contexts]

	self.assertEqual(shards, [(0, 4), (1, 5), (2, 6), (3, 7)])
	self.assertEqual(sorted(index for shard in shards for index in shard), list(range(8)))

	uneven = [
	assigned_group_indices(
	4,
	DistributedContext(rank=rank, local_rank=rank, world_size=3),
	)
	for rank in range(3)
	]
	self.assertEqual(uneven, [(0, 3), (1,), (2,)])

	def test_gradient_sum_handles_rank_local_and_globally_unused_parameters(self):
	with tempfile.TemporaryDirectory() as directory:
	rendezvous = str(Path(directory) / "rendezvous")
	mp.spawn(
	_gradient_worker,
	args=(rendezvous, directory),
	nprocs=2,
	join=True,
	)
	results = [
	torch.load(
	Path(directory) / f"rank-{rank}.pt",
	weights_only=True,
	)
	for rank in range(2)
	]

	for result in results:
	self.assertEqual(result["shared"].item(), 3.0)
	self.assertEqual(result["rank_zero_only"].item(), 4.0)
	self.assertIsNone(result["globally_unused"])

	def test_two_rank_actor_update_matches_single_worker_global_group(self):
	actor = torch.nn.Linear(1, 1, bias=False)
	actor.weight.data.zero_()
	optimizer = torch.optim.SGD(actor.parameters(), lr=0.1)
	scheduler = torch.optim.lr_scheduler.LambdaLR(
	optimizer, lambda _step: 1.0
	)
	config = RLConfig()
	config.rollout.group_size = 2
	config.reference.beta = 0.0
	expected = execute_actor_update(
	groups=[_group()],
	policy_scorer=_Scorer(actor),
	config=config,
	parameters=list(actor.parameters()),
	actor=actor,
	optimizer=optimizer,
	scheduler=scheduler,
	)

	with tempfile.TemporaryDirectory() as directory:
	mp.spawn(
	_engine_worker,
	args=(str(Path(directory) / "engine-rendezvous"), directory),
	nprocs=2,
	join=True,
	)
	results = [
	torch.load(
	Path(directory) / f"engine-rank-{rank}.pt",
	weights_only=True,
	)
	for rank in range(2)
	]

	for result in results:
	self.assertTrue(torch.equal(result["weight"], actor.weight.detach()))
	self.assertEqual(result["grad_norm"], expected.grad_norm)
	self.assertEqual(result["records"], 2)

	def test_rank_local_stage_error_is_reported_by_every_worker(self):
	with tempfile.TemporaryDirectory() as directory:
	mp.spawn(
	_stage_error_worker,
	args=(str(Path(directory) / "stage-rendezvous"), directory),
	nprocs=2,
	join=True,
	)
	messages = [
	(Path(directory) / f"stage-rank-{rank}.txt").read_text()
	for rank in range(2)
	]

	self.assertEqual(messages[0], messages[1])
	self.assertIn("distributed rollout failed", messages[0])
	self.assertIn("rank 1: RuntimeError: fixture", messages[0])


	if __name__ == "__main__":
	unittest.main()